"""Grid selection and reproducible uncertainty/commissioning scenarios."""
from itertools import product
import json
from pathlib import Path
import numpy as np
from scipy.optimize import minimize_scalar
from model import candidate, frequency_response, simulate, runup_force, floor_bump, spin


def load_config(path='inputs.json'):
    c=json.loads(Path(path).read_text(encoding='utf-8'))
    for key in ['mass_factors','cg_shift_m','stiffness_factors','damping_factors']:
        values=c['uncertainty'][key]
        if not isinstance(values,list) or not values or not np.all(np.isfinite(values)):raise ValueError('Uncertainty axes must be nonempty finite lists.')
    if not c['candidate_pair_stiffness_N_m'] or not c['candidate_damping_ratios']:raise ValueError('Candidate axes cannot be empty.')
    # Construction validates the geometry and all mechanical scalar parameters.
    candidate(c,c['candidate_pair_stiffness_N_m'][0],c['candidate_damping_ratios'][0])
    if len(c['operating_frequency_hz'])!=2 or not 0<c['operating_frequency_hz'][0]<c['operating_frequency_hz'][1]:raise ValueError('Invalid operating band.')
    if not np.isfinite(c['unbalance_kg_m']) or c['unbalance_kg_m']<=0:raise ValueError('Invalid unbalance.')
    if any(not np.isfinite(v) or v<=0 for key,v in c['limits'].items() if key!='group_load_min_N'):raise ValueError('Limits must be finite and positive.')
    if c['limits']['group_load_min_N']<0:raise ValueError('Minimum compression-only reaction cannot be negative.')
    for k in c['candidate_pair_stiffness_N_m']:
        for z in c['candidate_damping_ratios']:
            if not np.isfinite(k) or k<=0 or not np.isfinite(z) or z<=0:raise ValueError('Invalid grid candidate.')
    run=c['runup'];bump=c['bump']
    spin(0,run['maximum_hz'],run['ramp_s'],run['hold_s'])
    floor_bump(bump['amplitude_m'],bump['duration_s'])
    if not np.isfinite(run['settle_s']) or run['settle_s']<=0 or not np.isfinite(bump['record_s']) or bump['record_s']<=bump['duration_s']:raise ValueError('The transient records must include a positive settling interval.')
    for m,cg,kf,kb,cf in corners(c):candidate(c,100000,.1,m,cg,(kf,kb),cf)
    return c


def corners(c):
    u=c['uncertainty']
    return list(product(u['mass_factors'],u['cg_shift_m'],u['stiffness_factors'],u['stiffness_factors'],u['damping_factors']))


def harmonic_metrics(skid,c,frequencies=None):
    freqs=np.linspace(*c['operating_frequency_hz'],151) if frequencies is None else np.asarray(frequencies)
    response=frequency_response(skid,freqs)
    force=c['unbalance_kg_m']*(2*np.pi*freqs)**2
    reaction=abs(response['reaction'])*force[:,None]
    relative=abs(response['relative'])*force[:,None]
    acceleration=abs(response['sensor_acceleration'])*force
    q,static,load=skid.static()
    return {'frequency':freqs,'reaction':reaction,'relative':relative,'sensor_acceleration':acceleration,
            'operating_reaction_envelope_N':float(np.max(np.sum(reaction,axis=1))),
            'operating_sensor_acceleration_m_s2':float(np.max(acceleration)),
            'static_compression_mm':float(np.max(static)*1000),
            'decay_2percent_s':float(-np.log(.02)/skid.decay_rate()),
            'group_load_min_N':float(np.min(load[None,:]-reaction)),
            'group_load_max_N':float(np.max(load[None,:]+reaction)),
            'total_compression_mm':float(np.max(static[None,:]+relative)*1000)}


def failures(metrics,limits):
    failures=[]
    for key,value in limits.items():
        if key not in metrics:continue
        if (metrics[key]<value if key=='group_load_min_N' else metrics[key]>value):failures.append(key)
    return failures


def grid_search(c):
    rows=[];unc=corners(c)
    for k,z in product(c['candidate_pair_stiffness_N_m'],c['candidate_damping_ratios']):
        metrics=[harmonic_metrics(candidate(c,k,z,m,cg,(kf,kb),cf),c) for m,cg,kf,kb,cf in unc]
        keys=[key for key in c['limits'] if key in metrics[0]]
        envelope={key:float(min(r[key] for r in metrics) if key=='group_load_min_N' else max(r[key] for r in metrics)) for key in keys}
        bad=failures(envelope,c['limits'])
        rows.append(dict(stiffness_N_m=k,damping_ratio=z,accepted=not bad,rejected_by='; '.join(bad),**envelope))
    feasible=[r for r in rows if r['accepted']]
    if not feasible:raise RuntimeError('No grid candidate meets the harmonic/static/decay envelope. Inspect the rejected candidates.')
    # A finite grid choice, not a continuous/global optimum.
    selected=min(feasible,key=lambda r:(r['operating_reaction_envelope_N'],r['stiffness_N_m'],r['damping_ratio']))
    return rows,selected


def transient_metrics(skid,c,dt=.002,rtol=1e-8):
    run=c['runup'];end=2*run['ramp_s']+run['hold_s']+run['settle_s']
    t=np.linspace(0,end,round(end/dt)+1)
    a=simulate(skid,t,force=runup_force(c['unbalance_kg_m'],run['maximum_hz'],run['ramp_s'],run['hold_s']),rtol=rtol)
    bump=c['bump'];t2=np.linspace(0,bump['record_s'],round(bump['record_s']/dt)+1)
    b=simulate(skid,t2,base=floor_bump(bump['amplitude_m'],bump['duration_s']),rtol=rtol)
    _,static,load=skid.static()
    # Triangle inequality bounds every relative phase/time offset of the two disturbances.
    travel=np.max(abs(a['relative']),axis=0)+np.max(abs(b['relative']),axis=0)
    reactions=np.max(abs(a['reaction']),axis=0)+np.max(abs(b['reaction']),axis=0)
    allowance=1.01
    metrics={'transient_relative_motion_mm':float(max(travel)*1000*allowance),
             'total_compression_mm':float(max(static+travel*allowance)*1000),
             'group_load_min_N':float(min(load-reactions*allowance)),
             'group_load_max_N':float(max(load+reactions*allowance)),
             'runup_peak_sensor_acceleration_m_s2':float(max(abs(a['sensor_acceleration']))),
             'bump_peak_sensor_acceleration_m_s2':float(max(abs(b['sensor_acceleration']))),
             'final_relative_motion_mm':float(max(abs(a['relative'][-1]))*1000),
             'energy_residual_J':float(max(np.max(abs(a['energy_residual'])),np.max(abs(b['energy_residual']))))}
    return metrics,a,b


def uncertainty_transients(c,selected):
    rows=[]
    for index,(m,cg,kf,kb,cf) in enumerate(corners(c)):
        skid=candidate(c,selected['stiffness_N_m'],selected['damping_ratio'],m,cg,(kf,kb),cf)
        metric,_,_=transient_metrics(skid,c)
        harmonic=harmonic_metrics(skid,c,np.linspace(*c['operating_frequency_hz'],601))
        for key in ['static_compression_mm','operating_reaction_envelope_N','operating_sensor_acceleration_m_s2','decay_2percent_s']:metric[key]=harmonic[key]
        # Include harmonic load/travel as well as the combined transient bound.
        metric['total_compression_mm']=max(metric['total_compression_mm'],harmonic['total_compression_mm'])
        metric['group_load_min_N']=min(metric['group_load_min_N'],harmonic['group_load_min_N'])
        metric['group_load_max_N']=max(metric['group_load_max_N'],harmonic['group_load_max_N'])
        bad=failures(metric,c['limits'])
        rows.append(dict(case=index+1,mass_factor=m,cg_shift_m=cg,front_stiffness_factor=kf,rear_stiffness_factor=kb,damping_factor=cf,accepted=not bad,rejected_by='; '.join(bad),**metric))
    return rows


def local_peak_refinement(skid,c):
    f=np.linspace(*c['operating_frequency_hz'],601)
    def objective(x):return harmonic_metrics(skid,c,[x])['operating_reaction_envelope_N']
    y=np.array([objective(x) for x in f])
    possibilities=[(float(y[0]),float(f[0])),(float(y[-1]),float(f[-1]))]
    for i in range(1,len(f)-1):
        if y[i]>=y[i-1] and y[i]>=y[i+1]:
            res=minimize_scalar(lambda x:-objective(x),bounds=(f[i-1],f[i+1]),method='bounded',options={'xatol':1e-10})
            possibilities.append((-float(res.fun),float(res.x)))
    value,frequency=max(possibilities)
    return {'refined_peak_N':value,'frequency_hz':frequency,'grid_peak_N':float(max(y)),'grid_to_refined_N':float(value-max(y))}