"""Rest-to-rest jerk-limited trajectories and horizontal screw-axis duty sizing.

SI units throughout. Each phase is an exact cubic position polynomial. Motor
integrals use polynomial antiderivatives, split at force/power sign changes.
This is a reference calculation, not a drive command generator or safety control.
"""
from dataclasses import dataclass, asdict
import math
import numpy as np
from numpy.polynomial import Polynomial as Poly


def finite(value, name, positive=False, nonnegative=False):
    if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
        raise ValueError(f'{name} must be a finite number')
    if positive and value <= 0 or nonnegative and value < 0:
        raise ValueError(f'{name} has an invalid sign')
    return float(value)


@dataclass(frozen=True)
class Phase:
    start: float
    duration: float
    x: float
    v: float
    a: float
    jerk: float

    def state(self, local):
        q = float(local)
        return (self.x+self.v*q+self.a*q*q/2+self.jerk*q*q*q/6,
                self.v+self.a*q+self.jerk*q*q/2,
                self.a+self.jerk*q, self.jerk)


@dataclass(frozen=True)
class Profile:
    distance: float
    velocity_limit: float
    acceleration_limit: float
    jerk_limit: float
    phases: tuple
    duration: float
    regime: str

    def state(self, t):
        t = finite(t, 'time', nonnegative=True)
        if t > self.duration+1e-12:
            raise ValueError('time is outside the move')
        if not self.phases or t >= self.duration:
            return (self.distance, 0., 0., 0.)
        for phase in self.phases:
            if t < phase.start+phase.duration:
                return phase.state(t-phase.start)
        return (self.distance, 0., 0., 0.)

    def sample(self, dt=0.002):
        dt=finite(dt,'sample interval',positive=True)
        if self.duration/dt > 2_000_000:
            raise ValueError('sample request exceeds 2,000,000 intervals')
        if not self.phases:
            return np.array([[0.,0.,0.,0.,0.]])
        times=np.unique(np.r_[np.arange(0,self.duration,dt),
                              [p.start for p in self.phases],self.duration])
        return np.array([[t,*self.state(float(t))] for t in times])


def plan(distance, velocity, acceleration, jerk):
    """Symmetric seven-phase rest-to-rest solution; zero-length phases omitted."""
    distance=finite(distance,'distance')
    velocity=finite(velocity,'velocity limit',positive=True)
    acceleration=finite(acceleration,'acceleration limit',positive=True)
    jerk=finite(jerk,'jerk limit',positive=True)
    if (0<abs(distance)<1e-12) or abs(distance)>1000 or max(velocity,acceleration,jerk)>1e6 or min(velocity,acceleration,jerk)<1e-9:
        raise ValueError('input is outside the supported SI numerical range')
    if distance==0:
        return Profile(0.,velocity,acceleration,jerk,(),0.,'zero distance')
    distance_abs=abs(distance)
    tj=min(acceleration/jerk,math.sqrt(velocity/jerk))
    ta=max(0.,velocity/(jerk*tj)-tj)
    peak_velocity=jerk*tj*(tj+ta)
    needed=peak_velocity*(2*tj+ta)
    if distance_abs>=needed:
        cruise=(distance_abs-needed)/peak_velocity
        regime='velocity limited' if ta==0 else 'full seven-phase'
    else:
        cruise=0.
        tj=(distance_abs/(2*jerk))**(1/3)
        if jerk*tj<=acceleration:
            ta=0.
            regime='triangular acceleration'
        else:
            tj=acceleration/jerk
            ta=(-3*tj+math.sqrt(tj*tj+4*distance_abs/acceleration))/2
            regime='acceleration limited'
    durations=[tj,ta,tj,cruise,tj,ta,tj]
    signs=[1,0,-1,0,-1,0,1]
    direction=math.copysign(1,distance)
    phases=[];x=v=a=time=0.
    for dt, sign in zip(durations,signs):
        if dt<=1e-15:
            continue
        phase=Phase(time,dt,x,v,a,direction*sign*jerk)
        phases.append(phase)
        x,v,a,_=phase.state(dt)
        time+=dt
    if not math.isfinite(time):
        raise ValueError('non-finite trajectory duration')
    if abs(x-distance)>max(1e-10,1e-9*distance_abs) or abs(v)>1e-8 or abs(a)>1e-8:
        raise ArithmeticError('trajectory endpoint did not close')
    return Profile(distance,velocity,acceleration,jerk,tuple(phases),time,regime)


@dataclass(frozen=True)
class Axis:
    mass_kg: float=22.
    lead_m_rev: float=.02
    forward_efficiency: float=.9
    backdrive_efficiency: float=.85
    coulomb_force_N: float=28.
    linear_viscous_Ns_m: float=15.
    rotating_inertia_kgm2: float=.00008
    rotary_viscous_Nms: float=.00002
    torque_constant_Nm_Arms: float=.5
    phase_resistance_ohm: float=.7
    drive_motoring_efficiency: float=.95
    drive_regen_efficiency: float=.85

    def __post_init__(self):
        for name,value in asdict(self).items():
            finite(value,name,positive=name not in ('coulomb_force_N','linear_viscous_Ns_m','rotating_inertia_kgm2','rotary_viscous_Nms','phase_resistance_ohm'),
                   nonnegative=name in ('coulomb_force_N','linear_viscous_Ns_m','rotating_inertia_kgm2','rotary_viscous_Nms','phase_resistance_ohm'))
            if 'efficiency' in name and value>1:
                raise ValueError(f'{name} cannot exceed one')

    @property
    def radius(self):
        return self.lead_m_rev/(2*math.pi)

    def load(self,v,a,direction=0):
        """Return shaft torque and powers; friction is zero during a true dwell."""
        moving=abs(v)>1e-14 or abs(a)>1e-14
        direction=math.copysign(1,v) if abs(v)>1e-14 else direction
        force=self.mass_kg*a+(self.coulomb_force_N*direction if moving else 0)+self.linear_viscous_Ns_m*v
        linear=force*v
        factor=1/self.forward_efficiency if linear>=0 else self.backdrive_efficiency
        omega=v/self.radius
        torque=force*self.radius*factor+self.rotating_inertia_kgm2*a/self.radius+self.rotary_viscous_Nms*omega
        shaft=torque*omega
        copper=3*self.phase_resistance_ohm*(torque/self.torque_constant_Nm_Arms)**2
        terminal=shaft+copper
        bus=terminal/self.drive_motoring_efficiency if terminal>=0 else terminal*self.drive_regen_efficiency
        return {'torque_Nm':torque,'speed_rpm':omega*60/(2*math.pi),'current_Arms':torque/self.torque_constant_Nm_Arms,
                'shaft_power_W':shaft,'copper_loss_W':copper,'bus_power_W':bus,'linear_power_W':linear,
                'transmission_loss_W':force*self.radius*factor*omega-linear}


def real_roots(poly,low,high):
    roots=[]
    for root in poly.roots():
        if abs(np.imag(root))<1e-8 and low+1e-12<float(np.real(root))<high-1e-12:
            roots.append(float(np.real(root)))
    return sorted(set(roots))


def integral(poly,low,high):
    anti=poly.integ()
    return float(anti(high)-anti(low))


def positive_negative(poly,low,high):
    cuts=[low,*real_roots(poly,low,high),high];positive=negative=0.
    for a,b in zip(cuts,cuts[1:]):
        value=integral(poly,a,b)
        if poly((a+b)/2)>=0:positive+=value
        else:negative+=value
    return positive,negative


def peak_abs(poly,low,high):
    points=[low,high,*real_roots(poly.deriv(),low,high)]
    return max(abs(float(poly(t))) for t in points)


def move_metrics(profile,axis):
    totals=dict(duration_s=profile.duration,torque_square_integral_Nm2s=0.,peak_torque_Nm=0.,peak_speed_rpm=0.,
                shaft_net_J=0.,mechanical_regen_J=0.,bus_import_J=0.,bus_regen_J=0.,copper_J=0.,
                friction_J=0.,transmission_loss_J=0.,rotary_drag_J=0.)
    direction=math.copysign(1,profile.distance) if profile.distance else 0
    r=axis.radius
    for phase in profile.phases:
        v=Poly([phase.v,phase.a,phase.jerk/2]);a=Poly([phase.a,phase.jerk]);omega=v/r
        force=axis.mass_kg*a+axis.coulomb_force_N*direction+axis.linear_viscous_Ns_m*v
        cuts=[0.,*real_roots(force,0.,phase.duration),phase.duration]
        totals['peak_speed_rpm']=max(totals['peak_speed_rpm'],peak_abs(omega,0,phase.duration)*60/(2*math.pi))
        for left,right in zip(cuts,cuts[1:]):
            mid=(left+right)/2
            factor=1/axis.forward_efficiency if force(mid)*v(mid)>=0 else axis.backdrive_efficiency
            transmitted=force*r*factor
            torque=transmitted+axis.rotating_inertia_kgm2*a/r+axis.rotary_viscous_Nms*omega
            shaft=torque*omega
            copper=3*axis.phase_resistance_ohm*(torque/axis.torque_constant_Nm_Arms)**2
            friction=axis.coulomb_force_N*direction*v+axis.linear_viscous_Ns_m*v*v
            transmission=transmitted*omega-force*v
            drag=axis.rotary_viscous_Nms*omega*omega
            totals['torque_square_integral_Nm2s']+=integral(torque*torque,left,right)
            totals['peak_torque_Nm']=max(totals['peak_torque_Nm'],peak_abs(torque,left,right))
            totals['shaft_net_J']+=integral(shaft,left,right)
            totals['mechanical_regen_J']-=positive_negative(shaft,left,right)[1]
            positive,negative=positive_negative(shaft+copper,left,right)
            totals['bus_import_J']+=positive/axis.drive_motoring_efficiency
            totals['bus_regen_J']-=negative*axis.drive_regen_efficiency
            for key,poly in [('copper_J',copper),('friction_J',friction),('transmission_loss_J',transmission),('rotary_drag_J',drag)]:
                totals[key]+=integral(poly,left,right)
    totals['energy_balance_error_J']=totals['shaft_net_J']-totals['friction_J']-totals['transmission_loss_J']-totals['rotary_drag_J']
    totals['rms_torque_Nm']=math.sqrt(max(0.,totals['torque_square_integral_Nm2s']/profile.duration)) if profile.duration else 0.
    return totals


def cycle_metrics(distance,velocity,acceleration,jerk,axis,dwell_s):
    if len(dwell_s)!=2:
        raise ValueError('supply two dwell durations')
    dwell=[finite(d,'dwell',nonnegative=True) for d in dwell_s]
    outward=plan(distance,velocity,acceleration,jerk);homeward=plan(-distance,velocity,acceleration,jerk)
    metrics=[move_metrics(p,axis) for p in (outward,homeward)]
    result={k:(max(m[k] for m in metrics) if k.startswith('peak_') else sum(m[k] for m in metrics)) for k in metrics[0] if k!='rms_torque_Nm'}
    result['duration_s']+=sum(dwell)
    result['rms_torque_Nm']=math.sqrt(result['torque_square_integral_Nm2s']/result['duration_s']) if result['duration_s'] else 0.
    result['rms_current_Arms']=result['rms_torque_Nm']/axis.torque_constant_Nm_Arms
    result['peak_current_Arms']=result['peak_torque_Nm']/axis.torque_constant_Nm_Arms
    result['average_bus_regen_W']=result['bus_regen_J']/result['duration_s'] if result['duration_s'] else 0.
    result['cycles_per_minute']=60/result['duration_s'] if result['duration_s'] else 0.
    return result


def acceptance(metrics,limits):
    requirements={'peak_torque_Nm':'peak_torque_Nm','rms_torque_Nm':'continuous_torque_Nm',
                  'peak_current_Arms':'peak_current_Arms','rms_current_Arms':'continuous_current_Arms',
                  'peak_speed_rpm':'maximum_speed_rpm','bus_regen_J':'regen_energy_per_cycle_J',
                  'average_bus_regen_W':'average_regen_W'}
    checks={key:metrics[key]<=finite(limits[target],target,positive=True)*(1+1e-10) for key,target in requirements.items()}
    return all(checks.values()),checks