"""Dependency correctness and numerical isolation of generated local solves."""
from copy import deepcopy
from dataclasses import replace
import json
import math
import re
import subprocess
import unittest

from app.simulation.native_codegen.build import build_native, toolchain
from app.simulation.native_codegen.compiler import NativeCapabilityError
from app.simulation.native_codegen.extended import compile_extended_program
from app.simulation.native_codegen.schedule import Computation as Op, EvaluationSchedule
from tests.native_reference import reference_data, reference_network


class DependencyGraphTests(unittest.TestCase):
    def test_deep_reverse_emission_order_and_state_origin(self):
        ops = [Op.assignment(str(i), f'h[{i}]', f'h[{i-1}]', 'alias') for i in range(2000, 0, -1)]
        plan = EvaluationSchedule(ops, {'h[0]': 'state:chamber.m,U'})
        self.assertFalse(any(block.cyclic for block in plan.blocks))
        self.assertEqual([plan.computations[b.members[0]].key for b in plan.blocks], list(map(str, range(1, 2001))))
        self.assertEqual(plan.report()['blocks'][-1]['origins'], ['state:chamber.m,U'])

    def test_separate_cycles_and_dependent_output(self):
        ops = [Op.assignment('result', 'q[0]', 'h[1]+h[3]'),
               Op.assignment('a', 'h[1]', 'h[2]+g[0].h', 'stream'),
               Op.assignment('b', 'h[2]', 'h[1]', 'alias'),
               Op.assignment('c', 'h[3]', 'h[4]+g[1].h', 'stream'),
               Op.assignment('d', 'h[4]', 'h[3]', 'alias')]
        plan = EvaluationSchedule(ops, {'g[0].h': 'state:left', 'g[1].h': 'state:right'})
        self.assertEqual([b.members for b in plan.blocks if b.cyclic], [(1, 2), (3, 4)])
        self.assertEqual(plan.blocks[-1].members, (0,))
        self.assertEqual(plan.report()['blocks'][-1]['origins'], ['state:left', 'state:right'])

    def test_missing_source_duplicate_producer_and_source_free_alias_cycle(self):
        with self.assertRaisesRegex(NativeCapabilityError, 'missing native input'):
            EvaluationSchedule([Op.assignment('a', 'h[1]', 'h[0]')], {'h[1]': 'initial guess'})
        with self.assertRaisesRegex(NativeCapabilityError, 'Multiple native producers'):
            EvaluationSchedule([Op.assignment('a', 'q[0]', '0'), Op.assignment('b', 'q[0]', '1')], {})
        plan = EvaluationSchedule([Op.assignment('a', 'h[0]', 'h[1]', 'alias'),
                                   Op.assignment('b', 'h[1]', 'h[0]', 'alias')],
                                  {'h[0]': 'initial guess', 'h[1]': 'initial guess'})
        with self.assertRaisesRegex(NativeCapabilityError, 'no thermodynamic source'):
            plan.emit()


def prefix_case(case, prefix):
    result = deepcopy(case)
    for component in result['components']:
        component['name'] = prefix + component['name']
    result['connections'] = [[[prefix+name, port] for name, port in edge] for edge in result['connections']]
    return result


def instrument(program):
    """Count actual emitted operation executions; no timing counters in production."""
    count = program.evaluation_schedule['operationCount']
    code = re.sub(r'/\* schedule operation (\d+): \w+ \*/', r'++schedule_counts[\1];', program.source)
    code = code.replace('int model_eval(', 'static int model_eval_impl(', 1)
    code = '#include <stdio.h>\nstatic unsigned long long schedule_counts['+str(count)+'];\n'+code
    code += '''
int model_eval(double t,const double *y,double *dy,double *w) {
    for(int i=0;i<'''+str(count)+''';i++) schedule_counts[i]=0;
    int ok=model_eval_impl(t,y,dy,w);
    fprintf(stderr,"[");
    for(int i=0;i<'''+str(count)+''';i++) fprintf(stderr,"%s%llu",i?",":"",schedule_counts[i]);
    fprintf(stderr,"]\\n");return ok;
}
'''
    return replace(program, source=code)


class NativeScheduleTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        try:
            toolchain()
        except (OSError, RuntimeError, subprocess.SubprocessError) as exc:
            raise unittest.SkipTest(f'Native toolchain unavailable: {exc}')

    def probe(self, program, states):
        build = build_native(instrument(program))
        inputs = ''.join(' '.join(format(v, '.17g') for v in (0, *[state[k] for k in program.state_keys]))+'\n' for state in states)
        run = subprocess.run([str(build.executable), '--probe'], input=inputs, capture_output=True, text=True, check=True, timeout=30)
        rows = [json.loads(line) for line in run.stdout.splitlines()]
        counts = [json.loads(line) for line in run.stderr.splitlines()]
        self.assertEqual(len(rows), len(states))
        self.assertEqual(len(counts), len(states))
        for row in rows:
            self.assertTrue(row['success'])
        return rows, counts

    def test_two_pressure_loops_do_not_recompute_each_other_or_direct_branch(self):
        data = reference_data()['cases']
        a, b, direct = prefix_case(data[37], 'A_'), prefix_case(data[37], 'B_'), prefix_case(data[22], 'C_')
        combined = {'components': a['components']+b['components']+direct['components'],
                    'connections': a['connections']+b['connections']+direct['connections']}
        base = data[37]['probes'][0]
        bbase = data[37]['probes'][-1]
        # Independent regions use different states, then all reverse pressure direction.
        states = {}
        for prefix, case, probe in [('A_', data[37], base), ('B_', data[37], bbase), ('C_', data[22], data[22]['probes'][0])]:
            states.update({prefix+key: value for key, value in zip(case['stateKeys'], probe['state'])})
        reverse = dict(states)
        for prefix in ('A_', 'B_', 'C_'):
            for field in ('m', 'U'):
                left, right = prefix+'left.'+field, prefix+'right.'+field
                reverse[left], reverse[right] = states[right], states[left]
        inputs = [states, reverse, states]
        together = compile_extended_program(reference_network(combined))
        self.assertEqual(together.evaluation_schedule['cyclicBlockCount'], 2)
        actual, counts = self.probe(together, inputs)
        together_ops = together.evaluation_schedule['operations']
        for case in (a, b, direct):
            alone = compile_extended_program(reference_network(case))
            expected, alone_counts = self.probe(alone, inputs)
            for row, wanted in zip(actual, expected):
                output_map = dict(zip((v.key for v in together.variables), row['outputs']))
                state_map = dict(zip(together.state_keys, row['rhs']))
                self.assertEqual([output_map[v.key] for v in alone.variables], wanted['outputs'])
                self.assertEqual([state_map[key] for key in alone.state_keys], wanted['rhs'])
            for i, op in enumerate(alone.evaluation_schedule['operations']):
                if op['kind'] != 'flow':
                    continue
                j = next(j for j, candidate in enumerate(together_ops) if candidate['key']==op['key'])
                self.assertEqual([row[j] for row in counts], [row[i] for row in alone_counts])
                if case is direct:
                    self.assertEqual([row[j] for row in counts], [1, 1, 1])

    def test_reference_aliases_are_prepared_once_even_with_reverse_component_order(self):
        case = deepcopy(reference_data()['cases'][30])
        original = compile_extended_program(reference_network(case))
        case['components'].reverse()
        case['connections'].reverse()
        reordered = compile_extended_program(reference_network(case))
        inputs = [dict(zip(case['stateKeys'], probe['state'])) for probe in case['probes']]
        expected, _ = self.probe(original, inputs)
        actual, counts = self.probe(reordered, inputs)
        self.assertEqual(reordered.evaluation_schedule['cyclicBlockCount'], 0)
        self.assertTrue(all(value==1 for row in counts for value in row))
        for row, wanted in zip(actual, expected):
            outputs = dict(zip((v.key for v in reordered.variables), row['outputs']))
            rhs = dict(zip(reordered.state_keys, row['rhs']))
            # Permuting connection equations can change floating-point summation
            # order in the constant linear elimination, by roundoff only.
            for got, want in zip([outputs[v.key] for v in original.variables]+[rhs[key] for key in original.state_keys],
                                 wanted['outputs']+wanted['rhs']):
                self.assertTrue(math.isclose(got,want,rel_tol=5e-14,abs_tol=2e-12), (got,want))

    def test_coupled_pressure_and_mixing_loop_conserves_mass_and_energy(self):
        from tests.test_native_catalog import Circuit
        b = Circuit()
        node = b.add('tee','junction')
        for i,(port,pressure,temp) in enumerate((('port_in',4e5,330),('port_out1',1e5,280),('port_out2',2e5,300))):
            chamber = b.chamber('storage'+str(i),p0=pressure,T0=temp)
            pipe = b.add('amesim_pnl00r','pipe'+str(i),diam=.01,le=1,rr=1e-5)
            b.connect(chamber,'port_1',pipe,'port_1')
            b.connect(pipe,'port_2',node,port)
        program = compile_extended_program(b.seal())
        loops = [block for block in program.evaluation_schedule['blocks'] if block['cyclic']]
        self.assertEqual(len(loops),1)
        self.assertTrue(loops[0]['pressureUnknowns'])
        self.assertIn('stream:junction',loops[0]['operations'])
        build = build_native(program)
        initial = json.loads(subprocess.run([str(build.executable),'--init'],capture_output=True,text=True,check=True,timeout=15).stdout)
        state = dict(zip(program.state_keys,initial))
        reverse = dict(state)
        equal = dict(state)
        for field in ('m','U'):
            reverse['storage0.'+field],reverse['storage1.'+field] = state['storage1.'+field],state['storage0.'+field]
            for i in (1,2): equal['storage'+str(i)+'.'+field] = state['storage0.'+field]
        rows,_ = self.probe(program,[state,reverse,equal,state])
        self.assertEqual(rows[0],rows[-1])
        for row in rows:
            rhs = dict(zip(program.state_keys,row['rhs']))
            for field in ('m','U'):
                rates = [rhs['storage'+str(i)+'.'+field] for i in range(3)]
                self.assertLessEqual(abs(sum(rates)), 1e-10 + 1e-9*sum(map(abs,rates)))


if __name__ == '__main__':
    unittest.main()