"""Scientific invariants and six required smoke checks; never writes formal weights."""
from __future__ import annotations

import json
import time

import numpy as np
import torch

from common import (DarcyEnergy, analytical_coefficient_tensor, arguments, build_model, configuration,
                    element_corners, evaluate, initialize, relative_l2, read_data, output_path, write_json)


def gauss_energy(p, u, h):
    """Independent high-order quadrature, only for verification of the analytical loss."""
    nodes, weights = np.polynomial.legendre.leggauss(3)
    corners_p, corners_u = element_corners(p), element_corners(u)
    answer = torch.zeros(p.shape[0], dtype=torch.float64)
    for gx, wx in zip((nodes + 1)/2, weights/2):
        for gy, wy in zip((nodes + 1)/2, weights/2):
            shape = torch.tensor([(1-gx)*(1-gy), gx*(1-gy), gx*gy, (1-gx)*gy])
            dx = torch.tensor([-(1-gy), 1-gy, gy, -gy]) / h
            dy = torch.tensor([-(1-gx), -gx, gx, 1-gx]) / h
            coefficient = (corners_p * shape).sum(-1)
            value = 0.5 * coefficient * ((corners_u*dx).sum(-1)**2+(corners_u*dy).sum(-1)**2) - (corners_u*shape).sum(-1)
            answer += value.sum((1,2)) * h**2 * wx * wy
    return answer


def main():
    args = arguments(__doc__).parse_args()
    cfg = configuration(args)
    device = initialize(cfg)
    fields, metadata = read_data(cfg)
    checks = {}
    start = time.perf_counter()
    # Validate the energy implementation independently in float64.
    exact = DarcyEnergy(5).double()
    p = torch.rand(2,5,5,1,dtype=torch.float64)+1
    u = torch.rand_like(p)
    analytical = exact(u, p, reduction="none")
    quadrature = gauss_energy(p, u, 0.25)
    torch.testing.assert_close(analytical, quadrature, rtol=2e-7, atol=1e-8)
    constant_stiffness = analytical_coefficient_tensor(1,1).sum(0)
    expected = torch.tensor([[4,-1,-2,-1],[-1,4,-1,-2],[-2,-1,4,-1],[-1,-2,-1,4]],dtype=torch.float64)/6
    torch.testing.assert_close(constant_stiffness, expected)
    grad_u = torch.rand(1,3,3,1,dtype=torch.float64,requires_grad=True)
    torch.autograd.gradcheck(lambda z: DarcyEnergy(3).double()(z,torch.ones_like(z)), (grad_u,))
    checks["analytical_energy"] = {"status":"passed", "quadrature_max_difference": (analytical-quadrature).abs().max().item(), "constant_stiffness":"passed", "autograd_gradcheck":"passed"}
    model = build_model(cfg,device)
    objective = DarcyEnergy(64).to(device)
    coefficient = fields["P_train"][:2].to(device)
    output = model(coefficient)
    assert output.shape == coefficient.shape and torch.isfinite(output).all()
    assert output[:,0].abs().max() == 0 and output[:,-1].abs().max() == 0
    assert output[:,:,0].abs().max() == 0 and output[:,:,-1].abs().max() == 0
    checks["forward_pass"] = "passed"
    objective(output,coefficient).backward()
    assert all(param.grad is not None and torch.isfinite(param.grad).all() for param in model.parameters())
    checks["backward_pass"] = "passed"
    optimizer = torch.optim.Adam(model.parameters(),lr=cfg["training"]["learning_rate"])
    model.train()
    losses=[]
    for batch in fields["P_train"].split(cfg["training"]["batch_size"]):
        batch = batch.to(device)
        optimizer.zero_grad(set_to_none=True)
        loss = objective(model(batch), batch)
        assert torch.isfinite(loss)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    assert losses[-1] < losses[0]
    checks["train_loop_dry_run"] = {"status":"passed", "samples":1000, "first_energy":losses[0],"last_energy":losses[-1]}
    _, errors, _ = evaluate(model, fields["P_test"].to(device),fields["U_test"].to(device),objective,100)
    assert len(errors)==100 and torch.isfinite(errors).all()
    checks["validation_loop_dry_run"]="passed"
    reference=fields["U_test"][:3]
    torch.testing.assert_close(relative_l2(reference,reference),torch.zeros(3))
    torch.testing.assert_close(relative_l2(2*reference,reference),torch.ones(3))
    checks["domain_metric_dry_run"]="passed"
    assert cfg["training"]["epochs"]==500 and cfg["model"]["depth"]==4 and cfg["model"]["width"]==32 and cfg["model"]["modes"]==12
    checks["config_consistency"]="passed"
    report={"status":"passed","attempts":1,"checks":checks,"data":metadata,"seconds":time.perf_counter()-start,
            "device":str(device),"device_name":torch.cuda.get_device_name() if device.type=="cuda" else "cpu"}
    write_json(output_path(cfg["output"]["results_dir"])/"verification.json",report)
    print(json.dumps(report,indent=2),flush=True)


if __name__=="__main__":
    main()