#!/usr/bin/env python3
"""Conditional response inference for a synthetic equal-density Maxwell world.
Python 3.11+, standard library only. No Lab or private repository imports.
"""
import argparse
import csv
import hashlib
import json
import math
from pathlib import Path

G, R, DENSITY, PARENT_MASS, SEMIMAJOR_AXIS = 6.67430e-11, 500000.0, 1000.0, 5e26, 3e8
GRAVITY = 4 * math.pi * G * DENSITY * R / 3
ANGULAR_FREQUENCY = math.sqrt(G * PARENT_MASS / SEMIMAJOR_AXIS ** 3)


def geometric_stiffness(thickness_km):
    """Beuthe 1410.4735v1, Appendix A, equation 108; equal-density shell."""
    x = 1 - 1000 * thickness_km / R
    return 24 / 5 * (19 - 75*x**3 + 112*x**5 - 75*x**7 + 19*x**10) / (24 + 40*x**3 - 45*x**7 - 19*x**10)


def prediction(frequency_ratio, logarithmic_parameters):
    effective_stiffness, maxwell_time = map(math.exp, logarithmic_parameters)
    argument = 1j * ANGULAR_FREQUENCY * frequency_ratio * maxwell_time
    return 0.6 / (0.4 + effective_stiffness * argument / (1 + argument))


def residuals(rows, parameters):
    values = []
    for row in rows:
        estimate = prediction(row['frequency_ratio'], parameters)
        values.extend([(row['real_k2'] - estimate.real) / row['sigma_real'], (row['imaginary_k2'] - estimate.imag) / row['sigma_imaginary']])
    return values


def objective(rows, parameters):
    return sum(value * value for value in residuals(rows, parameters))


def fit(rows):
    parameters = min(([a, b] for a in [i / 2 for i in range(-4, 21)] for b in [i / 2 for i in range(10, 29)]), key=lambda p: objective(rows, p))
    for _ in range(100):
        current = residuals(rows, parameters)
        derivatives = []
        for axis in range(2):
            upper, lower = parameters.copy(), parameters.copy()
            upper[axis] += 1e-5; lower[axis] -= 1e-5
            derivatives.append([(a - b) / 2e-5 for a, b in zip(residuals(rows, upper), residuals(rows, lower))])
        aa = sum(v*v for v in derivatives[0]); bb = sum(v*v for v in derivatives[1]); ab = sum(a*b for a,b in zip(*derivatives))
        ar = sum(a*b for a,b in zip(derivatives[0], current)); br = sum(a*b for a,b in zip(derivatives[1], current))
        determinant = aa * bb - ab * ab
        delta = [(bb*ar-ab*br)/determinant, (aa*br-ab*ar)/determinant]
        old = objective(rows, parameters)
        step = 1.0
        while step > 1e-8:
            candidate = [a-step*b for a,b in zip(parameters, delta)]
            if objective(rows, candidate) < old:
                parameters = candidate
                break
            step /= 2
        if max(abs(step*d) for d in delta) < 1e-10:
            break
    return parameters


def fixed_stiffness_profile(rows, thickness_km, shear_modulus):
    effective = geometric_stiffness(thickness_km) * shear_modulus / (DENSITY * GRAVITY * R)
    first = math.log(effective)
    def derivative(time_log):
        total = 0.0
        for row in rows:
            x = ANGULAR_FREQUENCY * row['frequency_ratio'] * math.exp(time_log)
            estimate = prediction(row['frequency_ratio'], [first, time_log])
            # Differentiate k=.6(1+i x)/(.4+i(.4+A)x) with respect to log(tau).
            gradient = -0.6j * effective * x / (0.4 + 1j*(0.4+effective)*x)**2
            total -= 2*((row['real_k2']-estimate.real)*gradient.real/row['sigma_real']**2 +
                        (row['imaginary_k2']-estimate.imag)*gradient.imag/row['sigma_imaginary']**2)
        return total
    grid = [5 + 9*index/48 for index in range(49)]
    candidates = [grid[0], grid[-1]]
    for lower, upper in zip(grid, grid[1:]):
        if derivative(lower) <= 0 <= derivative(upper):
            for _ in range(60):
                middle = (lower+upper)/2
                if derivative(middle) < 0: lower = middle
                else: upper = middle
            candidates.append((lower+upper)/2)
    second = min(candidates, key=lambda value: objective(rows,[first,value]))
    if second in [grid[0],grid[-1]]:
        raise RuntimeError('Conditional profile reached the numerical relaxation-time boundary')
    return objective(rows, [first, second]), second


def read_rows(path):
    with path.open(newline='') as stream:
        return [{key: float(value) for key, value in row.items()} for row in csv.DictReader(stream)]


def write_csv(path, rows):
    with path.open('w', newline='') as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0])); writer.writeheader(); writer.writerows(rows)


def analyze(input_directory, output_directory):
    output_directory.mkdir(parents=True, exist_ok=True)
    rows = read_rows(input_directory / 'synthetic_complex_response.csv')
    parameters = fit(rows)
    effective_dimensionless, maxwell_time = map(math.exp, parameters)
    effective_modulus = effective_dimensionless * DENSITY * GRAVITY * R
    comparison = []
    for row in rows:
        estimate = prediction(row['frequency_ratio'], parameters)
        comparison.append({**row, 'predicted_real_k2': estimate.real, 'predicted_imaginary_k2': estimate.imag,
                           'standardized_real_residual': (row['real_k2']-estimate.real)/row['sigma_real'],
                           'standardized_imaginary_residual': (row['imaginary_k2']-estimate.imag)/row['sigma_imaginary']})
    family = []
    for thickness in [10, 20, 50, 100, 200, 400, 500]:
        z = geometric_stiffness(thickness)
        modulus = effective_modulus / z
        family.append({'shell_thickness_km': thickness, 'has_fluid_in_this_model': thickness < 500,
                       'geometric_stiffness': z, 'shear_modulus_pa': modulus, 'viscosity_pa_s': modulus*maxwell_time,
                       'chi_squared': objective(rows, parameters)})
    profile = []
    for index in range(1761):
        thickness = 10 + index / 4
        chi, time_log = fixed_stiffness_profile(rows, thickness, 3.5e9)
        profile.append({'shell_thickness_km': thickness, 'fixed_shear_modulus_pa': 3.5e9, 'chi_squared': float(format(chi, '.10g')),
                        'profiled_viscosity_pa_s': float(format(3.5e9*math.exp(time_log), '.10g'))})
    best = min(profile, key=lambda row: row['chi_squared'])
    accepted = [row['shell_thickness_km'] for row in profile if row['chi_squared'] - best['chi_squared'] <= 1]
    love = prediction(1, parameters)
    power = -10.5 * G * PARENT_MASS**2 * R**5 / SEMIMAJOR_AXIS**6 * ANGULAR_FREQUENCY * 0.005**2 * love.imag
    failure = read_rows(input_directory / 'nonpassive_failure.csv')
    rejected = any(row['imaginary_k2'] > 5*row['sigma_imaginary'] for row in failure)
    summary = {'schema': 'sandboxphysics-ac12-conditional-inference-v1',
               'data_kind': 'Frozen synthetic harmonic-probe responses; not Cassini or planetary measurements',
               'time_convention': 'real part of complex amplitude times exp(+i omega t); passive imaginary k2 <= 0',
               'units': {'length': 'm', 'mass': 'kg', 'time': 's', 'shear_modulus': 'Pa', 'viscosity': 'Pa s', 'power': 'W'},
               'input_sha256': {name: hashlib.sha256((input_directory/name).read_bytes()).hexdigest() for name in ['synthetic_complex_response.csv', 'nonpassive_failure.csv']},
               'world': {'radius_m': R, 'density_kg_m3': DENSITY, 'gravity_m_s2': GRAVITY, 'parent_mass_kg': PARENT_MASS, 'semimajor_axis_m': SEMIMAJOR_AXIS, 'orbital_frequency_rad_s': ANGULAR_FREQUENCY},
               'fit': {'z_times_shear_modulus_pa': effective_modulus, 'maxwell_time_s': maxwell_time, 'chi_squared': objective(rows, parameters), 'numbers': 6, 'fitted_effective_parameters': 2, 'residual_degrees_of_freedom': 4, 'physical_orbit_power_at_e_0_005_w': power},
               'conditional_profile': {'assumption': 'shear modulus fixed exactly at 3.5 GPa; all other world properties fixed', 'grid_step_km': 0.25, 'best_grid_thickness_km': best['shell_thickness_km'], 'delta_chi_squared_1_grid_interval_km': [min(accepted),max(accepted)], 'interpretation': 'Local one-parameter Gaussian profile approximation, conditional on the stated material and noise model'},
               'nonpassive_failure': {'rejected': rejected, 'reason': 'positive imaginary response is incompatible with passive Maxwell loss under the declared sign convention'},
               'identifiability': 'Every family row has identical complex response at every frequency; changing thickness rescales modulus and viscosity together. This uniform-density result is not a general theorem about real stratified moons.'}
    write_csv(output_directory/'fitted_response.csv', comparison)
    write_csv(output_directory/'shell_solid_family.csv', family)
    write_csv(output_directory/'conditional_thickness_profile.csv', profile)
    (output_directory/'summary.json').write_text(json.dumps(summary, indent=2, sort_keys=True)+'\n')
    if not rejected: raise RuntimeError('The incompatible response was not rejected')
    print(json.dumps(summary['fit'], indent=2)); print(json.dumps(summary['conditional_profile'], indent=2))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input', type=Path, default=Path(__file__).resolve().parent/'data')
    parser.add_argument('--output', type=Path, default=Path.cwd()/'rerun-output')
    args = parser.parse_args(); analyze(args.input, args.output)


if __name__ == '__main__': main()