# Public copy: internal compute identifiers and local environment paths were removed or generalized. Scientific values are retained. Paths may need adapting and some original run inputs are not included.

"""DC resistivity forward + inversion validation on the compute environment.

Synthetic dipole-dipole survey over a conductive block in a half-space.
Proves: mesh build, survey, forward simulation, noise, inverse assembly,
directive stack, convergence to target misfit.
"""
import json, sys, time
import numpy as np
from discretize import TensorMesh
from discretize.utils import active_from_xyz
import simpeg
from simpeg import maps, data, data_misfit, regularization, optimization
from simpeg import inverse_problem, directives, inversion
from simpeg.electromagnetics.static import resistivity as dc
from simpeg.electromagnetics.static.utils.static_utils import (
    generate_dcip_sources_line, apparent_resistivity_from_voltage,
)

t0 = time.time()
rng = np.random.default_rng(42)
log = {"versions": {"simpeg": simpeg.__version__, "numpy": np.__version__},
       "formulation": "Simulation2DCellCentered/Robin"}

# ---- Survey: dipole-dipole line, 20 m electrode spacing --------------------
survey = generate_dcip_sources_line(
    survey_type="dipole-dipole", data_type="volt", dimension_type="2D",
    end_points=np.r_[-400.0, 400.0], topo=0.0, num_rx_per_src=10, station_spacing=20.0,
)
survey = dc.Survey(survey)
log["n_sources"] = len(survey.source_list)
log["nD"] = int(survey.nD)

# ---- Mesh: core 10 m cells + geometric padding ------------------------------
dh, npad = 10.0, 8
hx = [(dh, npad, -1.4), (dh, 100), (dh, npad, 1.4)]
hz = [(dh, npad, -1.4), (dh, 30)]
mesh = TensorMesh([hx, hz], origin="CN")
mesh.origin = np.r_[mesh.origin[0], -mesh.h[1].sum()]
log["mesh"] = {"n_cells": int(mesh.n_cells), "shape": [int(n) for n in mesh.shape_cells],
               "extent_x": [float(mesh.nodes_x.min()), float(mesh.nodes_x.max())],
               "extent_z": [float(mesh.nodes_y.min()), float(mesh.nodes_y.max())]}

topo = np.c_[np.linspace(-1500, 1500, 61), np.zeros(61)]
active = active_from_xyz(mesh, topo)
n_active = int(active.sum())
log["n_active"] = n_active

# ---- True model: 0.01 S/m half-space + 0.1 S/m block ------------------------
sigma_bg, sigma_block = 1e-2, 1e-1
true_model = np.full(n_active, sigma_bg)
cc = mesh.cell_centers[active, :]
block = (cc[:, 0] > -100) & (cc[:, 0] < 100) & (cc[:, 1] > -120) & (cc[:, 1] < -40)
true_model[block] = sigma_block
log["n_block_cells"] = int(block.sum())

exp_map = maps.ExpMap(nP=n_active)
act_map = maps.InjectActiveCells(mesh, active, 1e-8)
model_map = act_map * exp_map          # log-conductivity -> full mesh
true_log = np.log(true_model)

# Simulation2DNodal + default Robin BC cannot assemble J in simpeg 0.24.0
# (getADeriv _MBC_sigma broadcast bug) -> cell-centered with Robin BC.
sim = dc.Simulation2DCellCentered(mesh, survey=survey, sigmaMap=model_map, bc_type="Robin")

# ---- Forward gate -----------------------------------------------------------
dpred = sim.dpred(true_log)
app_res = apparent_resistivity_from_voltage(survey, dpred)
log["forward"] = {
    "all_finite": bool(np.all(np.isfinite(dpred))),
    "n_data": int(dpred.size),
    "app_res_min": float(app_res.min()), "app_res_max": float(app_res.max()),
    "app_res_median": float(np.median(app_res)),
    "background_res_true": 1.0 / sigma_bg,
}

# ---- Noise + uncertainties --------------------------------------------------
rel_err, floor = 0.05, 1e-5
std = rel_err * np.abs(dpred) + floor
dobs = dpred + rng.normal(size=dpred.size) * std
dc_data = data.Data(survey, dobs=dobs, standard_deviation=std)

# ---- Inverse problem --------------------------------------------------------
m0 = np.log(np.full(n_active, sigma_bg))
dmis = data_misfit.L2DataMisfit(simulation=sim, data=dc_data)
reg = regularization.WeightedLeastSquares(
    mesh, active_cells=active, reference_model=m0, alpha_s=1e-3, alpha_x=1.0, alpha_y=1.0,
)
opt = optimization.InexactGaussNewton(maxIter=40, maxIterCG=50)
# Default tolF/tolX/tolG are scaled by the INITIAL objective, which a large
# beta0_ratio makes enormous -> the optimizer "converges" long before the
# target misfit. Disable them and let TargetMisfit govern termination.
opt.tolF = 0.0
opt.tolX = 0.0
opt.tolG = 0.0
opt.eps = 0.0
inv_prob = inverse_problem.BaseInvProblem(dmis, reg, opt)

# VERIFIED against simpeg 0.24.0 source: directives.TargetMisfit stops on
# phi_d < chifact * phi_d_star, and phi_d_star defaults to nD itself
# (directives.py: self._phi_d_star = nD) -- NOT nD/2, despite L2DataMisfit
# carrying a leading 0.5. Judging against nD/2 reports a converged run as a
# failure.
chifact = 1.0
target_phi_d = chifact * survey.nD
log["target_phi_d"] = float(target_phi_d)

# Canonical SimPEG directive order: sensitivity weights -> beta estimate ->
# schedule -> target -> save -> preconditioner LAST. UpdatePreconditioner
# placed before BetaEstimate_ByEig builds the Jacobi preconditioner at the
# wrong beta, and the CG solves then stall with |fc-fOld| == 0 exactly.
dlist = [
    directives.UpdateSensitivityWeights(every_iteration=False),
    directives.BetaEstimate_ByEig(beta0_ratio=10.0, seed=42),
    directives.BetaSchedule(coolingFactor=2.0, coolingRate=1),
    directives.TargetMisfit(chifact=1.0),
    directives.SaveOutputEveryIteration(save_txt=True, name="dc_inv_output"),
    directives.UpdatePreconditioner(),
]
inv = inversion.BaseInversion(inv_prob, directiveList=dlist)
rec_log = inv.run(m0)

# ---- Convergence + model recovery ------------------------------------------
rec_sigma = np.exp(rec_log)
log["inversion"] = {
    "n_iterations": int(opt.iter),
    "final_phi_d": float(inv_prob.phi_d),
    "final_phi_m": float(inv_prob.phi_m),
    "final_beta": float(inv_prob.beta),
    "achieved_chi": float(inv_prob.phi_d / survey.nD),
    "reached_target": bool(inv_prob.phi_d < target_phi_d),
}
log["recovery"] = {
    "true_block_sigma": sigma_block,
    "rec_sigma_in_block_median": float(np.median(rec_sigma[block])),
    "rec_sigma_outside_median": float(np.median(rec_sigma[~block])),
    "contrast_recovered": float(np.median(rec_sigma[block]) / np.median(rec_sigma[~block])),
    "true_contrast": sigma_block / sigma_bg,
}
log["wall_seconds"] = round(time.time() - t0, 1)

np.save("recovered_model.npy", rec_sigma)
np.save("true_model.npy", true_model)
np.save("dobs.npy", dobs)
np.save("std.npy", std)
np.save("dpred_true.npy", dpred)
np.save("active_cells.npy", active)
np.save("cell_centers_active.npy", cc)
with open("validation.json", "w") as fh:
    json.dump(log, fh, indent=2)
print(json.dumps(log, indent=2))