include("mpc_core.jl")
include("neural_core.jl")
include("neural_adaptation.jl")
include("reference_config.jl")
include("imm_anchor.jl")
using Serialization
function design_load()
 folder=joinpath(@__DIR__,"results_mpc")
 mats=[Matrix{Float64}(readdlm(joinpath(folder,f*".csv"),',')) for f in ["A","B","K","P","H_terminal"]]
 (A=mats[1],B=mats[2],K=mats[3],P=mats[4],H=mats[5])
end
function episode(profile,policy,seed,design,bundle,out;duration=20.,radius=.001,budget=.002)
 models=DPDynamics.load_model_ratio.([1.,1.5]);data=DPDynamics.MuJoCo.init_data.(models)
 ENCODER_TRUST_RADIUS[]=radius;GATE_BUDGET[]=budget
 obs=policy in (:oracle,) ? nothing : policy==:ekf ? IMMAnchorObserver(bundle;learned=false,ratios=[1.],biasaware=false) : policy==:bias_ekf ? IMMAnchorObserver(bundle;learned=false,ratios=[1.]) : policy==:physics_only ? IMMAnchorObserver(bundle;learned=false) : policy in (:mm_frozen,:mm_adapt,:mm_binding) ? IMMAnchorObserver(bundle) : NeuralObserver(bundle,policy)
 if obs isa IMMAnchorObserver
  obs.base.policy=policy in (:mm_frozen,:physics_only,:ekf,:bias_ekf) ? :neural_frozen : :neural_encoder
  ANCHOR_CONFIG[]=(bias_sd=deg2rad(.03),residual_weight=0.,radius=radius)
 end
 c=LinearMPC(design.A,design.B,design.K,design.P,design.H);steps=round(Int,duration/.02)
 noise=randn(MersenneTwister(seed),NY,steps+1).*DPDynamics.SENSOR_STD;rng=MersenneTwister(seed+100000)
 x=vcat(.02+.002randn(rng),deg2rad(2.)+deg2rad(.2)*randn(rng),deg2rad(-1.).+deg2rad(.2).*randn(rng,NP-1),zeros(NY))
 uold=0.;viol=0;recover=0;shifted=0;accepts=0;attempts=0;blocks=0;complete=true;failure="";e2=0.;a2=0.;tracking2=0.;maxrho=0.;maxstep=0.;done=0
 file=joinpath(out,"$(profile)_$(policy)_seed$(seed).csv");mkpath(out)
 open(file,"w") do io
  headers=vcat(["time"],["x$j" for j in 1:NX],["estimate$j" for j in 1:NX],["reference$j" for j in 1:NX],["y$j" for j in 1:NY],["mask$j" for j in 1:NY],["u","violation","qp_accepted","shifted_plan","lqr_recovery","attempted","accepted","step","coupling","reference_coupling","loss_before","loss_after","gate_rejections","radius","gate_budget"])
  println(io,join(headers,','))
  for k in 0:steps-1
   try
    time=k*.02;available=!(profile in (:dropout,:combined)&&12<=time<12.4);b=profile in (:bias,:combined) ? deg2rad(3)*clamp(time-8,0.,1.) : 0.
    y=OBS_MAP*x+noise[:,k+1];y[2]+=b;mask=vcat(true,fill(available,NP));y[.!mask].=0.
    diag=(attempted=false,accepted=false,joint_step=0.,coupling=0.,loss_before=NaN,loss_after=NaN,budget_rejections=0);reference=copy(x);refcoupling=0.
    xhat=if policy==:oracle
     copy(x)
    else
     k>0&&predict_neural!(obs,uold);result=correct_neural!(obs,y,mask);diag=result.diagnostics;reference=result.reference_mean;refcoupling=result.reference_coupling;result.mean
    end
    control=solve_mpc!(c,xhat);u=control.u;audit=zeros(3)
    mode=profile in (:mass,:combined)&&time>=8 ? 2 : 1;xn=DPDynamics.transition!(models[mode],data[mode],x,u;audit)
    violation=audit[3]>0;viol+=violation;recover+=control.source=="uncertified_lqr";shifted+=control.source=="shifted_nominal"
    accepts+=diag.accepted;attempts+=diag.attempted;blocks+=diag.budget_rejections;maxrho=max(maxrho,diag.coupling);maxstep=max(maxstep,diag.joint_step)
    e=OBS_MAP*(xhat-x);e2+=sum(abs2,e);a2+=sum(abs2,e[2:NY]);tracking2+=sum(abs2,(OBS_MAP*x)[2:NY]);done+=1
    println(io,join([time;x;xhat;reference;y;Float64.(mask);u;Int(violation);Int(control.qp_accepted);Int(control.source=="shifted_nominal");Int(control.source=="uncertified_lqr");Int(diag.attempted);Int(diag.accepted);diag.joint_step;diag.coupling;refcoupling;diag.loss_before;diag.loss_after;diag.budget_rejections;radius;budget],','))
    x=xn;uold=u
   catch err
    if err isa ErrorException&&occursin("MuJoCo numerical warning",sprint(showerror,err))
     complete=false;failure=replace(sprint(showerror,err),'\n'=>' ');break
    end
    rethrow()
   end
  end
 end
 (profile=profile,method=policy,seed=seed,complete=complete,steps=done,violating_intervals=viol,zero_violation=complete&&viol==0,lqr_recovery=recover,shifted_plans=shifted,accepted=accepts,attempted=attempts,gate_rejections=blocks,max_coupling=maxrho,max_step=maxstep,angle_estimation_rmse_deg=done>0 ? rad2deg(sqrt(a2/(done*NP))) : NaN,angle_tracking_rmse_deg=done>0 ? rad2deg(sqrt(tracking2/(done*NP))) : NaN,diagnostic_mixed_rmse=done>0 ? sqrt(e2/done) : NaN,failure=failure,radius=radius,budget=budget)
end
function save_summary(rows,path)
 open(path,"w") do io
  ks=collect(keys(first(rows)));println(io,join(string.(ks),','))
  for row in rows;println(io,join((replace(string(getproperty(row,k)),','=>';') for k in ks),','));end
 end
end
function main()
 BLAS.set_num_threads(1);design=design_load();bundle=read_bundle();phase=get(ENV,"STANDARD_PHASE","smoke")
 out=joinpath(@__DIR__,phase);rows=Any[]
 if phase=="smoke"
  for policy in (:oracle,:ekf,:physics_only,:neural_frozen,:neural_encoder,:mm_adapt)
   q=episode(:nominal,policy,8101,design,bundle,out;duration=1.);push!(rows,q);println(q);flush(stdout)
  end
 elseif phase=="validation"
  for radius in (0.,.0001,.001,.01),seed in (8101,8102),profile in (:nominal,:combined)
   folder=joinpath(out,"radius_$(radius)");q=episode(profile,:neural_encoder,seed,design,bundle,folder;radius=radius);push!(rows,q);println(q);flush(stdout);save_summary(rows,joinpath(out,"radius_summary.csv"))
  end
 elseif phase=="calibration"
  radius=parse(Float64,read(joinpath(@__DIR__,"selected_radius.txt"),String))
  for profile in (:nominal,:combined)
   q=episode(profile,:mm_adapt,8103,design,bundle,out;radius);push!(rows,q);println(q);flush(stdout)
  end
  rho=Float64[]
  for profile in (:nominal,:combined)
   path=joinpath(out,"$(profile)_mm_adapt_seed8103.csv");names=split(readline(path),',');table=readdlm(path,',';skipstart=1)
   ci=findfirst(==("coupling"),names);ai=findfirst(==("accepted"),names)
   append!(rho,[Float64(table[k,ci]) for k in axes(table,1) if table[k,ai]>0&&table[k,ci]>0])
  end
  budget=median(rho)/2;@assert budget>1e-10
  serialize(joinpath(@__DIR__,"selected_config.jls"),(radius=radius,budget=budget))
  open(joinpath(@__DIR__,"selected_config.txt"),"w") do io;println(io,"radius=",radius,"\nbudget=",budget,"\nselection_seed=8103");end
  println("Frozen configuration ",(radius=radius,budget=budget));save_summary(rows,joinpath(out,"summary.csv"))
 else
  config=deserialize(joinpath(@__DIR__,"selected_config.jls"));seeds=(96301,96302,96303,96304,96305)
  policies=(:oracle,:ekf,:bias_ekf,:physics_only,:neural_gaussian,:neural_frozen,:neural_encoder,:neural_joint,:neural_no_trust,:mm_frozen,:mm_adapt,:mm_binding)
  for seed in seeds,profile in (:nominal,:mass,:bias,:dropout,:combined),policy in policies
   q=episode(profile,policy,seed,design,bundle,out;radius=config.radius,budget=policy==:mm_binding ? config.budget : .002);push!(rows,q);println(q);flush(stdout);save_summary(rows,joinpath(out,"summary.csv"))
  end
 end
end
abspath(PROGRAM_FILE)==abspath(@__FILE__)&&main()
