from difflib import unified_diff
import filecmp
import glob
import hashlib
from optparse import OptionParser
import os
import shutil
import sys

import numpy as np
import openmc
from openmc.examples import pwr_core
from colorama import Fore, init

from tests.regression_tests import config

init()


def colorize(diff):
    """Produce colored diff for test results"""
    for line in diff:
        if line.startswith('+'):
            yield Fore.RED + line + Fore.RESET
        elif line.startswith('-'):
            yield Fore.GREEN + line + Fore.RESET
        elif line.startswith('^'):
            yield Fore.BLUE + line + Fore.RESET
        else:
            yield line


class TestHarness:
    """General class for running OpenMC regression tests."""

    def __init__(self, statepoint_name):
        self._sp_name = statepoint_name

    def main(self):
        """Accept commandline arguments and either run or update tests."""
        if config['update']:
            self.update_results()
        else:
            self.execute_test()

    def execute_test(self):
        """Run OpenMC with the appropriate arguments and check the outputs."""
        try:
            self._run_openmc()
            self._test_output_created()
            results = self._get_results()
            self._write_results(results)
            self._compare_results()
        finally:
            self._cleanup()

    def update_results(self):
        """Update the results_true using the current version of OpenMC."""
        try:
            self._run_openmc()
            self._test_output_created()
            results = self._get_results()
            self._write_results(results)
            self._overwrite_results()
        finally:
            self._cleanup()

    def _run_openmc(self):
        if config['mpi']:
            mpi_args = [config['mpiexec'], '-n', config['mpi_np']]
            openmc.run(openmc_exec=config['exe'], mpi_args=mpi_args,
              event_based=config['event'])
        else:
            openmc.run(openmc_exec=config['exe'], event_based=config['event'])

    def _test_output_created(self):
        """Make sure statepoint.* and tallies.out have been created."""
        statepoint = glob.glob(self._sp_name)
        assert len(statepoint) == 1, 'Either multiple or no statepoint files' \
            ' exist.'
        assert statepoint[0].endswith('h5'), \
            'Statepoint file is not a HDF5 file.'
        if os.path.exists('tallies.xml'):
            assert os.path.exists('tallies.out'), \
                'Tally output file does not exist.'

    def _get_results(self, hash_output=False):
        """Digest info in the statepoint and return as a string."""
        # Read the statepoint file.
        statepoint = glob.glob(self._sp_name)[0]
        with openmc.StatePoint(statepoint) as sp:
            outstr = ''
            if sp.run_mode == 'eigenvalue':
                # Write out k-combined.
                outstr += 'k-combined:\n'
                form = '{0:12.6E} {1:12.6E}\n'
                outstr += form.format(sp.k_combined.n, sp.k_combined.s)

            # Write out tally data.
            for i, tally_ind in enumerate(sp.tallies):
                tally = sp.tallies[tally_ind]
                results = np.zeros((tally.sum.size * 2, ))
                results[0::2] = tally.sum.ravel()
                results[1::2] = tally.sum_sq.ravel()
                results = ['{0:12.6E}'.format(x) for x in results]

                outstr += 'tally {}:\n'.format(i + 1)
                outstr += '\n'.join(results) + '\n'

        # Hash the results if necessary.
        if hash_output:
            sha512 = hashlib.sha512()
            sha512.update(outstr.encode('utf-8'))
            outstr = sha512.hexdigest()

        return outstr

    def _write_results(self, results_string):
        """Write the results to an ASCII file."""
        with open('results_test.dat', 'w') as fh:
            fh.write(results_string)

    def _overwrite_results(self):
        """Overwrite the results_true with the results_test."""
        shutil.copyfile('results_test.dat', 'results_true.dat')

    def _compare_results(self):
        """Make sure the current results agree with the reference."""
        compare = filecmp.cmp('results_test.dat', 'results_true.dat')
        if not compare:
            expected = open('results_true.dat').readlines()
            actual = open('results_test.dat').readlines()
            diff = unified_diff(expected, actual, 'results_true.dat',
                                'results_test.dat')
            print('Result differences:')
            print(''.join(colorize(diff)))
            os.rename('results_test.dat', 'results_error.dat')
        assert compare, 'Results do not agree'

    def _cleanup(self):
        """Delete statepoints, tally, and test files."""
        output = glob.glob('statepoint.*.h5')