import json

import numpy as np

from ..core.log import log
from ..core.system import System

_LCSIM_DEFAULT_SETTINGS = dict(
    L=10,  # Thickness (micrometer)
    K1=10.3,  # Frank constants (pN)
    K2=7.4,
    K3=1.6 * 10.3,
    gamma=203.6,  # (mPa s)
    omega0=5,  # radius of the Gaussian beam (micrometer)
    beamx=0.0,  # Center of Gaussian beam (micrometer)
    beamy=0.0,
    p0=-15,  # helix pitch if photosiomer is absent
    pinf=1.5,  # helix pitch if normal isomer is absent
    sz=20,  # Number of nodes per thickness
    sx=200,  # Number of nodes per width
    tau=1000,  # photoisomer decay time (s)
    I0toIs=1000,  # Beam intensity constant
    diso=20,  # Isotropic diffusion rate (micrometer^2 / s = 1e-12 m^2/s)
    dani=0,  # Anisotropic diffusion rate
    IbtoIs=0,
    backomega=5,
    backpar=5,
    voltage=0,  # Voltage peak to peak (sinusoidal). (Volt) V_rms=voltage/2/sqrt(2)
    epsilona=1,  # Dielectric anisotropy (dimensionless).
    fixed_dir=(0, 0, 1),
    beam_radial_index=0,  # Hyperbolic momentum charge
    beam_azimuthal_index=0,  # Indicate the orbital angular momentum of the beam
)


def load_magnes_npz(filename, path="PATH") -> tuple[np.ndarray, System]:
    raise NotImplementedError


def save_npz_lcsim(file, nn, settings, periodic=False):
    n = np.ascontiguousarray(nn, dtype=np.float32)
    if n.ndim != 5 or n.shape[-1] != 3:
        raise ValueError(f"`nn` must have shape (sx, sy, sz, rank, 3), got {n.shape}.")

    alpha = np.zeros(shape=n.shape[:3], dtype=np.float32)
    default = {**_LCSIM_DEFAULT_SETTINGS, "bc": (0, 1 if periodic else 0, 0)}
    s = {
        **default,
        **settings,
        "sz": n.shape[2],
        "sy": n.shape[1],
        "sx": n.shape[0],
    }

    print(f"{n.shape=} {alpha.shape=}", s)

    np.savez_compressed(
        file,
        PATH=n[None],
        photo_concentration=alpha,
        settings=json.dumps(s),
    )


def load_lcsim_npz(filename, path="PATH") -> tuple[np.ndarray, System, dict]:
    log.debug(f"Open '{filename}'")
    datafile = np.load(filename, allow_pickle=True)
    log.debug(f"Available keys {*datafile.keys(),}")
    images = datafile[path]
    if images.ndim != 6 or images.shape[-1] != 3:
        raise ValueError(f"Shape {images.shape} of '{path}' is not supported.")
    settings = json.loads(datafile["settings"][()]) if "settings" in datafile else None
    system = System.cubic(size=images.shape[1:4])
    return list(images), system, settings


__all__ = [
    "load_magnes_npz",
    "load_lcsim_npz",
    "save_npz_lcsim",
]
