"""Vertical/pitch isolation model. SI units, small motions about static equilibrium."""
from dataclasses import dataclass, replace
import math
import numpy as np
from scipy.integrate import solve_ivp
from scipy.linalg import eigh

G = 9.80665

@dataclass(frozen=True)
class Skid:
    mass: float
    inertia: float
    x: tuple
    stiffness: tuple
    damping: tuple
    force_x: float = 0.22
    sensor_x: float = 0.25

    def __post_init__(self):
        scalars = [self.mass, self.inertia, self.force_x, self.sensor_x]
        if not np.all(np.isfinite(scalars)) or self.mass <= 0 or self.inertia <= 0:
            raise ValueError('Mass and inertia must be positive and all coordinates finite.')
        for values, label in [(self.x,'positions'),(self.stiffness,'stiffness'),(self.damping,'damping')]:
            if len(values) != 2 or not np.all(np.isfinite(values)):
                raise ValueError('Exactly two finite support-group '+label+' values are required.')
        if min(self.stiffness) <= 0 or min(self.damping) < 0:
            raise ValueError('Stiffness must be positive and damping nonnegative.')
        if not self.x[0] < 0 < self.x[1]:
            raise ValueError('The center of mass must lie strictly between the support groups.')

    @property
    def B(self):
        return np.array([[1., self.x[0]], [1., self.x[1]]])

    @property
    def M(self):
        return np.diag([self.mass,self.inertia])

    @property
    def K(self):
        return self.B.T @ np.diag(self.stiffness) @ self.B

    @property
    def C(self):
        return self.B.T @ np.diag(self.damping) @ self.B

    def static(self):
        """Positive downward displacement and compression under gravity."""
        q=np.linalg.solve(self.K,[self.mass*G,0.])
        compression=self.B@q
        return q,compression,np.array(self.stiffness)*compression

    def modes(self):
        values,vectors=eigh(self.K,self.M)
        return np.sqrt(values)/(2*np.pi),vectors

    def decay_rate(self):
        A=np.block([[np.zeros((2,2)),np.eye(2)],[-np.linalg.solve(self.M,self.K),-np.linalg.solve(self.M,self.C)]])
        return float(-np.max(np.linalg.eigvals(A).real))


def frequency_response(skid, frequencies, force=1., base=0.):
    """Complex amplitude using Re(Q exp(i omega t)); base translation is common to both feet."""
    f=np.atleast_1d(np.asarray(frequencies,dtype=float))
    if not np.all(np.isfinite(f)) or np.min(f)<0 or not np.isfinite(force) or not np.isfinite(base):
        raise ValueError('Frequencies must be finite and nonnegative; amplitudes finite.')
    w=2*np.pi*f
    D=skid.K[None,:,:]-w[:,None,None]**2*skid.M+1j*w[:,None,None]*skid.C
    k=np.asarray(skid.stiffness); c=np.asarray(skid.damping)
    rhs=np.broadcast_to(np.array([1.,skid.force_x])*force,(len(f),2)).astype(complex).copy()
    rhs+=np.einsum('ij,nj->ni',skid.B.T,(k[None,:]+1j*w[:,None]*c)*base)
    q=np.linalg.solve(D,rhs[...,None])[...,0]
    relative=q@skid.B.T-base
    reaction=(k[None,:]+1j*w[:,None]*c)*relative
    sensor=-(w**2)*(q[:,0]+skid.sensor_x*q[:,1])
    return {'q':q,'relative':relative,'reaction':reaction,'sensor_acceleration':sensor}


def stiffness_derivative(skid, frequency, group, force=1.):
    """dQ/dk at fixed damping and forcing, from differentiating D Q = F."""
    if group not in (0,1):raise ValueError('Group must be 0 or 1.')
    if not np.isfinite(frequency) or frequency<0:raise ValueError('Invalid frequency.')
    w=2*np.pi*frequency
    D=skid.K-w*w*skid.M+1j*w*skid.C
    q=np.linalg.solve(D,np.array([1.,skid.force_x])*force)
    b=skid.B[group]
    return -np.linalg.solve(D,np.outer(b,b)@q)


def spin(t, maximum_hz=30., ramp_s=2., hold_s=1.):
    """Smooth frequency ramps and their exact phase integral; returns theta, omega, alpha."""
    if not np.all(np.isfinite([t,maximum_hz,ramp_s,hold_s])) or maximum_hz<=0 or ramp_s<=0 or hold_s<0:raise ValueError('Invalid spin schedule.')
    T=ramp_s;H=hold_s;F=maximum_hz;W=2*np.pi*F
    if t<=0:return 0.,0.,0.
    if t<T:
        s=t/T;return W*T*(s**3-.5*s**4), W*(3*s*s-2*s**3), W*(6*s-6*s*s)/T
    theta0=W*T*.5
    if t<T+H:return theta0+W*(t-T),W,0.
    if t<2*T+H:
        s=(t-T-H)/T
        return theta0+W*H+W*T*(s-s**3+.5*s**4),W*(1-3*s*s+2*s**3),-W*(6*s-6*s*s)/T
    return W*(T+H),0.,0.


def runup_force(unbalance,maximum_hz=30.,ramp_s=2.,hold_s=1.):
    if not np.isfinite(unbalance) or unbalance<0:raise ValueError('Unbalance must be finite and nonnegative.')
    def force(t):
        theta,omega,alpha=spin(t,maximum_hz,ramp_s,hold_s)
        return unbalance*(omega*omega*np.sin(theta)-alpha*np.cos(theta))
    return force


def floor_bump(amplitude=.001,duration=.15):
    if not np.isfinite(amplitude) or not np.isfinite(duration) or duration<=0:raise ValueError('Invalid bump.')
    def base(t):
        if t<=0 or t>=duration:return 0.,0.
        a=np.pi*t/duration
        return amplitude*np.sin(a)**2,amplitude*np.pi/duration*np.sin(2*a)
    return base


def simulate(skid,times,force=None,base=None,initial=None,rtol=1e-8,max_step=.003):
    """Integrate q, qdot, work and dissipated energy. Returns sampled mechanical states."""
    t=np.asarray(times,dtype=float)
    if t.ndim!=1 or len(t)<2 or not np.all(np.isfinite(t)) or np.any(np.diff(t)<=0) or t[0]<0:
        raise ValueError('Times must be a strictly increasing nonnegative vector.')
    if not np.isfinite(rtol) or rtol<=0 or not np.isfinite(max_step) or max_step<=0:raise ValueError('Invalid solver controls.')
    y0=np.zeros(6)
    if initial is not None:
        initial=np.asarray(initial,dtype=float)
        if initial.shape!=(4,) or not np.all(np.isfinite(initial)):raise ValueError('Initial state must have four finite values.')
        y0[:4]=initial
    force=force or (lambda _:0.)
    base=base or (lambda _:(0.,0.))
    B=skid.B;K=skid.K;C=skid.C;M=skid.M
    k=np.asarray(skid.stiffness);c=np.asarray(skid.damping);invM=np.diag(1/np.diag(M))
    fdir=np.array([1.,skid.force_x]);ones=np.ones(2)
    def rhs(time,state):
        q=state[:2];v=state[2:4];yb,ybd=base(time);F=force(time)
        applied=fdir*F+B.T@(k*yb+c*ybd)
        a=invM@(applied-C@v-K@q)
        relative=B@q-yb;relative_v=B@v-ybd
        reaction=k*relative+c*relative_v
        # Dynamic spring energy uses relative motion; foundation work is -R*ydot.
        power=F*(fdir@v)-ybd*np.sum(reaction)
        loss=np.dot(c*relative_v,relative_v)
        return np.r_[v,a,power,loss]
    sol=solve_ivp(rhs,(t[0],t[-1]),y0,t_eval=t,method='DOP853',rtol=rtol,atol=rtol*1e-3,max_step=max_step)
    if not sol.success:raise RuntimeError(sol.message)
    q=sol.y[:2].T;v=sol.y[2:4].T
    accelerations=np.array([rhs(time,state)[2:4] for time,state in zip(t,sol.y.T)])
    y=np.array([base(time)[0] for time in t]);yd=np.array([base(time)[1] for time in t])
    relative=q@B.T-y[:,None];relative_v=v@B.T-yd[:,None]
    reaction=relative*k+relative_v*c
    energy=.5*np.einsum('ni,ij,nj->n',v,M,v)+.5*np.sum(relative**2*k,axis=1)
    residual=energy-energy[0]-sol.y[4]+sol.y[5]
    sensor_a=accelerations[:,0]+skid.sensor_x*accelerations[:,1]
    return {'time':t,'q':q,'velocity':v,'acceleration':accelerations,'relative':relative,'reaction':reaction,'sensor_acceleration':sensor_a,'energy':energy,'work':sol.y[4],'dissipation':sol.y[5],'energy_residual':residual,'base':y,'force':np.array([force(time) for time in t]),'nfev':sol.nfev}


def candidate(config,stiffness,damping_ratio,mass_factor=1.,cg_shift=0.,k_factors=(1.,1.),c_factor=1.):
    mass=config['mass_kg']*mass_factor
    xs=tuple(x-cg_shift for x in config['support_positions_m'])
    k=np.asarray(config['stiffness_split'])*stiffness*np.asarray(k_factors)
    nominal_x=np.array(config['support_positions_m'])
    shares=np.array([nominal_x[1],-nominal_x[0]])/(nominal_x[1]-nominal_x[0])*config['mass_kg']
    nominal_k=np.asarray(config['stiffness_split'])*stiffness
    c=2*damping_ratio*np.sqrt(nominal_k*shares)*c_factor
    return Skid(mass,config['pitch_inertia_kg_m2']*mass_factor,xs,tuple(k),tuple(c),config['force_position_m']-cg_shift,config['sensor_position_m']-cg_shift)