"""Independent first-order HLL evaluation solver for the SAME dCDB9 case.

This module is imported only by result.py after PINN training and inference.
It cannot supply labels to train.py. Cell-centered 0.04 m mesh, CFL=0.25,
periodic boundary conditions; explicit unsplit forward Euler time stepping.
"""
import time

import numpy as np
import torch


@torch.no_grad()
def solve(config, x, y, times, path, device):
    started = time.monotonic()
    ny, nx = len(y), len(x)
    xy = torch.stack(torch.meshgrid(torch.as_tensor(y, dtype=torch.float32, device=device),
                                    torch.as_tensor(x, dtype=torch.float32, device=device), indexing='ij'))
    physics = config['physics']
    inside = xy.square().sum(dim=0) <= physics['dam_radius'] ** 2
    conserved = torch.zeros((3, ny, nx), dtype=torch.float32, device=device)
    conserved[0] = torch.where(inside, float(physics['depth_inside']), float(physics['depth_outside']))
    dx, dy = float(x[1] - x[0]), float(y[1] - y[0])
    gravity, cfl = physics['gravity'], config['evaluation']['hll_cfl']
    snapshots, diagnostics = [], []
    t, step = 0.0, 0
    initial_mass = conserved[0].double().sum().item() * dx * dy
    # Array axes [component,y,x]: x-momentum is component 1 / spatial dim 2;
    # y-momentum is component 2 / spatial dim 1.
    for target in times:
        while t < float(target) - 1e-12:
            h = conserved[0]
            sound = torch.sqrt(gravity * h)
            max_x = (conserved[1].abs() / h + sound).max().item()
            max_y = (conserved[2].abs() / h + sound).max().item()
            dt = min(cfl / (max_x / dx + max_y / dy), float(target) - t)
            fx = interface_flux(conserved, gravity, momentum=1, dim=2)
            fy = interface_flux(conserved, gravity, momentum=2, dim=1)
            conserved = conserved - (dt / dx) * (fx - torch.roll(fx, 1, 2)) - (dt / dy) * (fy - torch.roll(fy, 1, 1))
            t += dt
            step += 1
            if step % 100 == 0:
                print('HLL_PROGRESS', {'step': step, 't': t, 'dt': dt,
                      'min_depth': conserved[0].min().item()}, flush=True)
        if not torch.isfinite(conserved).all() or conserved[0].min() <= 0:
            raise FloatingPointError('HLL reference is nonfinite or violates wet-cell positivity.')
        primitive = torch.stack((conserved[0], conserved[1] / conserved[0], conserved[2] / conserved[0]), dim=-1)
        snapshots.append(primitive.cpu().numpy())
        mass = conserved[0].double().sum().item() * dx * dy
        diagnostics.append({'time': float(target), 'mass': mass,
                            'relative_mass_drift': (mass - initial_mass) / initial_mass,
                            'min_depth': conserved[0].min().item(), 'steps': step})
        print('HLL_SNAPSHOT', diagnostics[-1], flush=True)
    fields = np.stack(snapshots)
    np.savez_compressed(path, x=x, y=y, times=times, fields=fields,
                        variables=np.array(['h', 'u', 'v']), dx=dx, dy=dy, cfl=cfl)
    return fields, {'scheme': 'independent first-order unsplit HLL with Davis wave speeds',
                    'initial_condition': 'paper Eq.55', 'training_use': False,
                    'dx': dx, 'dy': dy, 'cfl': cfl, 'boundary': 'periodic',
                    'dtype': 'float32', 'steps': step, 'diagnostics': diagnostics,
                    'elapsed_seconds': time.monotonic() - started}


def interface_flux(state, gravity, momentum, dim):
    """HLL numerical flux on forward cell interfaces, with periodic wrap."""
    neighbor = torch.roll(state, -1, dim)
    hl, hr = state[0], neighbor[0]
    vl, vr = state[momentum] / hl, neighbor[momentum] / hr
    sl = torch.minimum(vl - (gravity * hl).sqrt(), vr - (gravity * hr).sqrt())
    sr = torch.maximum(vl + (gravity * hl).sqrt(), vr + (gravity * hr).sqrt())
    fl, fr = state * vl, neighbor * vr
    fl[momentum] += 0.5 * gravity * hl.square()
    fr[momentum] += 0.5 * gravity * hr.square()
    hll = (sr * fl - sl * fr + sl * sr * (neighbor - state)) / (sr - sl)
    return torch.where((sl >= 0)[None], fl, torch.where((sr <= 0)[None], fr, hll))