import math
import unittest
import numpy as np
from scipy.integrate import quad
from planner import Axis, plan, move_metrics, cycle_metrics, acceptance


class TrajectoryTests(unittest.TestCase):
    def test_triangular_closed_form(self):
        p=plan(.001,.8,4,20)
        self.assertEqual(p.regime,'triangular acceleration')
        self.assertAlmostEqual(p.duration,4*(.001/40)**(1/3),places=13)

    def test_acceleration_limited(self):
        p=plan(.2,5,1,10)
        self.assertEqual(p.regime,'acceleration limited')
        self.assertAlmostEqual(p.duration,math.sqrt(.1**2+4*.2)+.1,places=12)

    def test_velocity_limited(self):
        p=plan(1,.1,5,10)
        self.assertEqual(p.regime,'velocity limited')
        self.assertAlmostEqual(p.duration,10+2*math.sqrt(.1/10),places=12)

    def test_full_profile(self):
        self.assertEqual(len(plan(1,.5,1,10).phases),7)

    def test_zero(self):
        p=plan(0,.5,1,10)
        self.assertEqual(p.state(0),(0,0,0,0))
        self.assertEqual(len(p.sample()),1)

    def test_signed_symmetry(self):
        a=plan(.31,.7,3,40);b=plan(-.31,.7,3,40)
        self.assertEqual(a.duration,b.duration)
        for t in np.linspace(0,a.duration,31):
            np.testing.assert_allclose(a.state(float(t)), -np.array(b.state(float(t))),atol=1e-12)

    def test_breakpoint_continuity(self):
        p=plan(1,.5,1,10)
        for first,second in zip(p.phases,p.phases[1:]):
            np.testing.assert_allclose(first.state(first.duration)[:3],second.state(0)[:3],atol=1e-13)

    def test_parameter_grid_limits(self):
        for d in [1e-8,.001,.02,.2,2,-.17]:
            for v,a,j in [(1,2,10),(.1,5,10),(.45,1.6,48)]:
                p=plan(d,v,a,j);rows=p.sample(max(p.duration/300,.00001))
                self.assertLessEqual(np.max(abs(rows[:,2])),v+1e-10)
                self.assertLessEqual(np.max(abs(rows[:,3])),a+1e-10)
                self.assertLessEqual(np.max(abs(rows[:,4])),j+1e-10)
                self.assertAlmostEqual(p.phases[-1].state(p.phases[-1].duration)[0],d,places=10)

    def test_derivatives(self):
        p=plan(.32,.35,1.6,24)
        for phase in p.phases:
            h=phase.duration*1e-4;t=phase.duration*.45
            left=phase.state(t-h);right=phase.state(t+h);center=phase.state(t)
            self.assertAlmostEqual((right[0]-left[0])/(2*h),center[1],delta=1e-8)
            self.assertAlmostEqual((right[1]-left[1])/(2*h),center[2],delta=1e-8)
            self.assertAlmostEqual((right[2]-left[2])/(2*h),center[3],delta=1e-8)

    def test_invalid_limits(self):
        for values in [(1,0,1,1),(1,1,-1,1),(1,1,1,float('nan')),(float('inf'),1,1,1),(True,1,1,1),(1e-300,1,1,1)]:
            with self.assertRaises(ValueError):plan(*values)

    def test_time_bounds_and_resource_limit(self):
        p=plan(.32,.35,1.6,24)
        with self.assertRaises(ValueError):p.state(-.1)
        with self.assertRaises(ValueError):p.state(p.duration+1)
        with self.assertRaises(ValueError):p.sample(1e-10)


class SizingTests(unittest.TestCase):
    def test_efficiency_sign(self):
        axis=Axis(rotating_inertia_kgm2=0,rotary_viscous_Nms=0,coulomb_force_N=0,linear_viscous_Ns_m=0)
        for v,a in [(1,1),(1,-1),(-1,-1),(-1,1)]:
            load=axis.load(v,a)
            factor=1/.9 if v*a>0 else .85
            self.assertAlmostEqual(load['torque_Nm'],axis.mass_kg*a*axis.radius*factor)
            self.assertGreaterEqual(load['transmission_loss_W'],-1e-12)

    def test_lossless_energy_balance(self):
        axis=Axis(forward_efficiency=1,backdrive_efficiency=1,coulomb_force_N=0,linear_viscous_Ns_m=0,rotary_viscous_Nms=0,phase_resistance_ohm=0,drive_motoring_efficiency=1,drive_regen_efficiency=1)
        p=plan(.32,.35,1.6,24);m=move_metrics(p,axis)
        self.assertAlmostEqual(m['shaft_net_J'],0,places=11)
        self.assertAlmostEqual(m['bus_import_J'],m['bus_regen_J'],places=11)
        kinetic=.5*(axis.mass_kg+axis.rotating_inertia_kgm2/axis.radius**2)*.35**2
        self.assertAlmostEqual(m['mechanical_regen_J'],kinetic,places=11)

    def test_dissipative_energy_balance(self):
        for d in [.001,.32,-.32]:
            m=move_metrics(plan(d,.35,1.6,24),Axis())
            self.assertLess(abs(m['energy_balance_error_J']),1e-10)
            self.assertGreater(m['friction_J'],0)
            self.assertGreater(m['transmission_loss_J'],0)

    def test_quadrature_independent_integral(self):
        axis=Axis();p=plan(.32,.35,1.6,24);m=move_metrics(p,axis)
        independent=0.
        for phase in p.phases:
            def integrand(q):
                _,v,a,_=phase.state(q)
                return axis.load(v,a,1)['torque_Nm']**2
            independent+=quad(integrand,0,phase.duration,epsabs=1e-11,epsrel=1e-10,limit=100)[0]
        self.assertAlmostEqual(m['torque_square_integral_Nm2s'],independent,delta=1e-9)

    def test_peak_matches_dense_sample(self):
        axis=Axis();p=plan(.32,.35,1.6,24);m=move_metrics(p,axis)
        sampled=max(abs(axis.load(row[2],row[3],1)['torque_Nm']) for row in p.sample(.0001))
        self.assertGreaterEqual(m['peak_torque_Nm']+1e-12,sampled)
        self.assertLess(m['peak_torque_Nm']-sampled,1e-5)

    def test_signed_duty_symmetry(self):
        a=move_metrics(plan(.32,.35,1.6,24),Axis());b=move_metrics(plan(-.32,.35,1.6,24),Axis())
        for key in a:self.assertAlmostEqual(a[key],b[key],delta=1e-10)

    def test_dwell_changes_rms_not_energy(self):
        a=cycle_metrics(.32,.35,1.6,24,Axis(),[0,0]);b=cycle_metrics(.32,.35,1.6,24,Axis(),[1,1])
        self.assertAlmostEqual(a['bus_regen_J'],b['bus_regen_J'])
        self.assertLess(b['rms_torque_Nm'],a['rms_torque_Nm'])

    def test_invalid_axis(self):
        for params in [{'mass_kg':-1},{'forward_efficiency':1.1},{'phase_resistance_ohm':-1},{'lead_m_rev':0},{'mass_kg':float('nan')}]:
            with self.assertRaises(ValueError):Axis(**params)

    def test_zero_cycle(self):
        m=cycle_metrics(0,.35,1.6,24,Axis(),[0,0])
        self.assertEqual(m['duration_s'],0)
        self.assertEqual(m['rms_torque_Nm'],0)

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