import numpy as np
from scipy.integrate import solve_ivp
import argparse
import json
import csv
from pathlib import Path
import time

def simulate_deployment(k_spring, theta_toggle, v_egress, I_fin, Cm_alpha, timer_latency_ms):
    """
    Simplified dynamical model for fin deployment.
    I_fin * theta'' = k_spring - M_aero - damping * theta'
    """
    rho = 1.225
    S_fin = 0.0036
    c_fin = 0.060
    
    # Aerodynamic hinge moment vs angle (simplified)
    # M_aero max when deployed
    def M_aero(theta):
        return 0.5 * rho * (v_egress**2) * S_fin * c_fin * Cm_alpha * np.sin(theta)
        
    damping = 0.05
    
    def dynamics(t, state):
        theta, omega = state
        
        # Delay start by latency
        if t < latency_s:
            return [0, 0]
            
        # Stop at toggle
        if theta >= theta_toggle:
            return [0, 0]
            
        torque_spring = k_spring # Constant force spring assumption
        torque_aero = M_aero(theta)
        
        alpha = (torque_spring - torque_aero - damping * omega) / I_fin
        return [omega, alpha]

    # Time span: wait for latency, then deploy
    latency_s = timer_latency_ms / 1000.0
    
    # Actually, latency just shifts the start time, doesn't affect final state unless rocket slows down.
    # We just simulate the deployment phase.
    
    # Initial state
    y0 = [0.0, 0.0]
    
    # We want to find the state at t_end (e.g., 0.2s)
    # Event: reaching toggle
    def hit_toggle(t, y):
        return y[0] - theta_toggle
    hit_toggle.terminal = True
    
    sol = solve_ivp(dynamics, [0, 0.2], y0, events=hit_toggle, max_step=0.01)
    
    theta_final = sol.y[0, -1]
    
    # If it hit the toggle event (status == 1), the mechanism mechanically locks,
    # so the final velocity is 0. Otherwise, it's the velocity at t_end.
    if sol.status == 1:
        omega_final = 0.0
    else:
        omega_final = sol.y[1, -1]
    
    # Convert to deg and deg/s
    theta_deg = np.degrees(theta_final)
    omega_deg_s = np.degrees(omega_final)
    
    return theta_deg, omega_deg_s

def run_monte_carlo(num_trials=10000, seed=33, output_dir=None):
    np.random.seed(seed)
    
    # Baselines
    k_base = 220.0
    theta_toggle_base = np.radians(90.0)
    v_base = 70.0
    I_base = (1/3) * 0.015 * (0.060**2)
    Cm_base = 0.8
    lat_base = 50.0
    
    # Sample distributions
    k_samples = np.random.uniform(k_base*0.9, k_base*1.1, num_trials)
    theta_samples = np.random.uniform(theta_toggle_base - np.radians(5.0), theta_toggle_base + np.radians(5.0), num_trials)
    v_samples = np.random.uniform(v_base*0.9, v_base*1.1, num_trials)
    I_samples = np.random.uniform(I_base*0.85, I_base*1.15, num_trials)
    Cm_samples = np.random.uniform(Cm_base*0.75, Cm_base*1.25, num_trials)
    lat_samples = np.random.uniform(lat_base - 10, lat_base + 10, num_trials)
    
    successes = 0
    failures = []
    
    for i in range(num_trials):
        th, om = simulate_deployment(k_samples[i], theta_samples[i], v_samples[i], I_samples[i], Cm_samples[i], lat_samples[i])
        
        # Success criteria: It must lock (reach its specific toggle angle) and come to a rest (v <= 50)
        toggle_deg = np.degrees(theta_samples[i])
        if abs(th - toggle_deg) <= 0.1 and abs(om) <= 50.0:
            successes += 1
        else:
            failures.append({
                'trial': i,
                'theta_final': th,
                'omega_final': om,
                'v_egress': v_samples[i],
                'Cm_alpha': Cm_samples[i]
            })
            
    success_rate = successes / num_trials
    
    if output_dir:
        csv_path = Path(output_dir) / "C7_reliability.csv"
        csv_path.parent.mkdir(parents=True, exist_ok=True)
        with open(csv_path, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=['trial', 'theta_final', 'omega_final', 'v_egress', 'Cm_alpha'],
                                    lineterminator='\n')
            writer.writeheader()
            writer.writerows(failures)

        # The failure-only CSV cannot carry its own denominator; this record does.
        record = {
            'schema': 'project33.c7.run_record/1',
            'generator': 'Simulation/reliability_sweep.py',
            'trials': num_trials,
            'successes': successes,
            'failures': len(failures),
            'seed': seed,
            'rng': 'numpy.random.seed + numpy.random.uniform (legacy global RNG)',
            'sampled_ranges': {
                'k_spring': [k_base * 0.9, k_base * 1.1],
                'theta_toggle_deg': [90.0 - 5.0, 90.0 + 5.0],
                'v_egress_m_s': [v_base * 0.9, v_base * 1.1],
                'I_fin_kg_m2': [I_base * 0.85, I_base * 1.15],
                'Cm_alpha': [Cm_base * 0.75, Cm_base * 1.25],
                'timer_latency_ms': [lat_base - 10, lat_base + 10],
            },
            'fixed_parameters': {'S_fin_m2': 0.0036, 'c_fin_m': 0.060, 'fin_mass_kg': 0.015,
                                 'fin_length_m': 0.060, 'damping_N_m_s': 0.05, 'rho_kg_m3': 1.225},
            'success_criterion': '|theta_final - theta_toggle| <= 0.1 deg and |omega_final| <= 50 deg/s',
            'model_structure_notes': [
                'omega_final is set to 0 whenever the toggle event fires, so the omega test cannot fail for a locked trial',
                'k_spring (stated in N*m/rad) is applied as a constant torque in N*m',
                'timer_latency_ms delays the start but has no effect on the final state',
                'fin mass, size and egress speed differ from the C6 design point',
            ],
            'success_fraction': successes / num_trials if num_trials else None,
            'basis': 'numerical simulation; no physical deployment was observed',
        }
        (csv_path.parent / "C7_run_record.json").write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8")

    return success_rate, failures

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--emit', required=True)
    args = parser.parse_args()
    
    print("Running Monte Carlo simulation (10,000 trials)...")
    t0 = time.time()
    sr, fails = run_monte_carlo(num_trials=10000, output_dir=Path(args.emit).parent)
    print(f"Finished in {time.time()-t0:.2f}s. Success Rate: {sr*100:.2f}%")