# turbo_wrapper.py
import math
import os
from typing import Literal

import numpy as np
# --- ONLY: SciPy FFT threading control (CPU) ---
import scipy.fft as spfft

from cupystorm import turbo_simulator as dns_all


class DnsSimulator:
    """
    DnsSimulator side is assumed to expose:
      - dns_init(N, re, k0)
      - dns_step(t, dt, cn)
      - dns_kinetic(px, py)
      - dns_om2phys(px, py)
      - dns_streamfunc(px, py)
      - dns_set_ur_real(field, nx, ny)
    """

    # Variable selector constants (used by the Qt GUI)
    VAR_U = 0
    VAR_V = 1
    VAR_ENERGY = 2
    VAR_OMEGA = 3
    VAR_STREAM = 4

    def __init__(
            self,
            n: int = 128,
            re: float = 100000,
            k0: float = 10.0,
            cfl: float = 0.25,
            backend: Literal["cpu", "gpu", "auto"] = "auto",
            mode: Literal["pao", "highh", "rain", "circle", "mouse"] = "pao",
            seed: int = 1,
    ):
        self.N = int(n)
        self.m = 3 * self.N
        self.re = float(re)
        self.k0 = float(k0)
        self.cfl = float(cfl)
        self.seed = int(seed)
        self.backend = str(backend)
        self.mode = str(mode)
        self.max_steps = 5000

        # --- ONLY: max SciPy FFT workers on CPU ---
        self.fft_workers = 4

        # UR dimensions from Fortran workspace: UR(2+3N/2, 3N/2, 3)
        # For the pure-Python solver, we use the full 3/2-grid from DnsState.
        #   ur_full has shape (3, NZ_full, NX_full)
        #   NZ_full = 3*N/2, NX_full = 3*N/2
        # We map these directly to (py, px) for the GUI.
        with spfft.set_workers(self.fft_workers):
            self.state = dns_all.create_dns_state(
                N=self.N,
                Re=self.re,
                K0=self.k0,
                CFL=self.cfl,
                backend=self.backend,  # GUI uses the CPU/NumPy/GPU/CuPy backend
                mode=self.mode,
                seed=self.seed,
            )

            self.nx = int(self.state.NZ_full)  # "height"
            self.ny = int(self.state.NX_full)  # "width"

            # expose for GUI sizing
            self.py = self.nx  # height
            self.px = self.ny  # width

            # time integration scalars
            self.t = float(self.state.t)
            self.dt = float(self.state.dt)
            self.cn = float(self.state.cn)
            self.iteration = 0

            # which field to visualize
            self.current_var = self.VAR_U

            # initialize Python DNS state (mirror dns_all.run_dns NEXTDT INIT)
            #   1) initial STEP2A from spectral to physical
            #   2) compute CFLM
            #   3) set DT and CN from CFL condition
            dns_all.dns_step2a(self.state)
            # CFLM * DT * PI = CFLNUM  →  DT = CFLNUM / (CFLM * PI)
            CFLM = dns_all.compute_cflm(self.state)
            CFLM0 = float(CFLM)  # <-- forces CPU scalar on GPU backend
            if CFLM0 == 0.0:
                self.state.dt = 0.01
            else:
                self.state.dt = float(self.state.cflnum) / (CFLM0 * math.pi)
            self.state.cn = 1.0
            self.state.cnm1 = 0.0

            # mirror into the simulator scalars
            self.t = float(self.state.t)
            self.dt = float(self.state.dt)
            self.cn = float(self.state.cn)

    # ------------------------------------------------------------------
    def step(self, mod_next_dt: int) -> None:
        """Advance one DNS step on the Fortran side."""
        # In the pure-Python version this mirrors dns_all.run_dns:
        #   dt_old = DT
        #   STEP2B
        #   STEP3
        #   STEP2A
        #   NEXTDT
        #   T = T + dt_old
        S = self.state

        if S.backend == "gpu":
            import cupy as cp  # type: ignore
            if isinstance(S.dt, cp.ndarray):
                dt_old = float(S.dt.item())
            elif isinstance(S.dt, cp.generic):
                dt_old = float(S.dt.item())
            else:
                dt_old = float(S.dt)
        else:
            dt_old = float(S.dt)


        if S.backend == "gpu":
            # Keep DT as a host float on GPU so STEP3 never sees a CuPy scalar/0-d array.
            S.dt = dt_old
        if S.backend == "cpu":
            with spfft.set_workers(self.fft_workers):
                dns_all.dns_step2b(S)
                dns_all.dns_step3(S)
                dns_all.dns_step2a(S)
                dns_all.next_dt(S)
        else:
            dns_all.dns_step2b(S)
            dns_all.dns_step3(S)
            dns_all.dns_step2a(S)
            if (self.iteration < 100) or (self.iteration % mod_next_dt) == 0:
                dns_all.next_dt(S)

        S.t = float(S.t) + dt_old
        self.t = float(S.t)
        self.dt = float(S.dt)
        self.cn = float(S.cn)
        self.iteration += 1

    def set_N(self, N: int) -> None:
        """Recreate the entire DNS state with a new grid size N."""
        self.N = int(N)
        self.m = 3 * self.N  # preserve original structure

        # Rebuild state exactly the same way __init__ does
        with spfft.set_workers(self.fft_workers):
            self.state = dns_all.create_dns_state(
                N=self.N,
                Re=self.re,
                K0=self.k0,
                CFL=self.cfl,
                backend=self.backend,
                mode=self.mode,
                seed=self.seed,
            )

        self.nx = int(self.state.NZ_full)
        self.ny = int(self.state.NX_full)
        self.py = self.nx
        self.px = self.ny

        # Reset integrator scalars like in reset_field()
        if self.state.backend == "cpu":
            with spfft.set_workers(self.fft_workers):
                dns_all.dns_step2a(self.state)
                CFLM = dns_all.compute_cflm(self.state)
                CFLM0 = float(CFLM)  # <-- forces CPU scalar on GPU backend
        else:
            dns_all.dns_step2a(self.state)
            CFLM = dns_all.compute_cflm(self.state)
            CFLM0 = float(CFLM)  # <-- forces CPU scalar on GPU backend

        if CFLM0 == 0.0:
            self.state.dt = 0.01
        else:
            self.state.dt = float(self.state.cflnum) / (CFLM0 * math.pi)
        self.state.cn = 1.0
        self.state.cnm1 = 0.0

        self.t = float(self.state.t)
        self.dt = float(self.state.dt)
        self.cn = float(self.state.cn)
        self.iteration = 0

    # ------------------------------------------------------------------
    def reset_field(self, mode: str | None = None) -> None:
        """Reinitialize the DNS state on the Fortran side."""

        self.t = np.float32(0.0)
        self.dt = np.float32(0.0)
        self.cn = np.float32(1.0)
        self.iteration = 0
        if mode is not None:
            self.mode = str(mode)

        # Pick a fresh PAO seed each reset (LCG is mod 5011 → use 1..5010)
        seed = 1 + (int.from_bytes(os.urandom(8), "little") % 5010)

        # Recreate the Python DNS state and redo the NEXTDT INIT phase
        with spfft.set_workers(self.fft_workers):
            self.state = dns_all.create_dns_state(
                N=self.N,
                Re=self.re,
                K0=self.k0,
                CFL=self.cfl,
                backend=self.backend,
                mode=self.mode,
                seed=seed,
            )

        self.nx = int(self.state.NZ_full)
        self.ny = int(self.state.NX_full)
        self.py = self.nx
        self.px = self.ny

        if self.state.backend == "cpu":
            with spfft.set_workers(self.fft_workers):
                dns_all.dns_step2a(self.state)
                CFLM = dns_all.compute_cflm(self.state)
                CFLM0 = float(CFLM)  # <-- forces CPU scalar on GPU backend
        else:
            dns_all.dns_step2a(self.state)
            CFLM = dns_all.compute_cflm(self.state)
            CFLM0 = float(CFLM)  # <-- forces CPU scalar on GPU backend

        if CFLM0 == 0.0:
            self.state.dt = 0.01
        else:
            self.state.dt = float(self.state.cflnum) / (CFLM0 * math.pi)
        self.state.cn = 1.0
        self.state.cnm1 = 0.0

        self.t = float(self.state.t)
        self.dt = float(self.state.dt)
        self.cn = float(self.state.cn)

    def set_body_force(
            self,
            ix_full: int,
            iz_full: int,
            amp: float = 0.2,
            sigma: float = 4.0,
            active: bool = True,
    ) -> None:
        S = self.state
        S.force_active = bool(active)
        S.force_dirty = True
        S.force_ix_full = int(ix_full)
        S.force_iz_full = int(iz_full)
        S.force_amp_x = float(amp)
        S.force_amp_z = 0.0
        S.force_sigma = float(sigma)

    # ------------------------------------------------------------------
    @staticmethod
    def _float_to_pixels(field: np.ndarray) -> np.ndarray:
        """
        Map a float field to 8-bit grayscale [1,255], like the PGM dumper.
        NaN/Inf are mapped to mid-gray.
        """
        fmin = float(np.nanmin(field))
        fmax = float(np.nanmax(field))
        rng = fmax - fmin

        if (not np.isfinite(rng)) or abs(rng) <= 1.0e-12:
            return np.full(field.shape, 128, dtype=np.uint8)

        norm = (field - fmin) / rng  # 0..1 (but may contain NaN/Inf)
        norm = np.nan_to_num(norm, nan=0.5, posinf=1.0, neginf=0.0)

        pixf = 1.0 + norm * 254.0  # 1..255
        pix = np.clip(pixf, 1.0, 255.0)

        return pix.astype(np.uint8)

    # ------------------------------------------------------------------
    def _snapshot(self, comp: int) -> np.ndarray:
        """
        Raw snapshot from Fortran, now using dns_frame with 3× scale-up.
        """
        # For the pure-Python solver, we take the component 'comp' from ur_full:
        #   comp = 1,2,3 → ur_full[0,1,2]
        S = self.state

        idx = int(comp) - 1
        if idx < 0 or idx > 2:
            idx = 0

        if S.backend == "gpu":
            import cupy as cp  # type: ignore
            field = cp.asnumpy(S.ur_full[idx, :, :])
        else:
            field = np.asarray(S.ur_full[idx, :, :])

        return self._float_to_pixels(field)

    # ------------------------------------------------------------------
    def make_pixels(self, comp: int = 1) -> np.ndarray:
        """
        Convenience: comp-based visualization.
        comp = 1,2,3 (UR components).
        """
        return self._snapshot(comp)

    # ------------------------------------------------------------------
    def make_pixels_component(self, var: int | None = None) -> np.ndarray:
        """
        High-level selector used by the GUI:
          VAR_U      -> UR(:,:,1)
          VAR_V      -> UR(:,:,2)
          VAR_ENERGY -> sqrt(u^2+v^2+w^2)
          VAR_OMEGA  -> vorticity
          VAR_STREAM -> streamfunction
        """
        if var is None:
            var = self.current_var

        S = self.state

        if var == self.VAR_U:
            plane = self._snapshot(1)
        elif var == self.VAR_V:
            plane = self._snapshot(2)
        elif var == self.VAR_ENERGY:
            # Use dns_all kinetic helper: fills ur_full[2,:,:]
            dns_all.dns_kinetic(S)
            if S.backend == "gpu":
                import cupy as cp  # type: ignore
                field = cp.asnumpy(S.ur_full[2, :, :])
            else:
                field = np.asarray(S.ur_full[2, :, :])
            plane = self._float_to_pixels(field)
        elif var == self.VAR_OMEGA:
            # Use dns_all omega→physical helper: fills ur_full[2,:,:]
            dns_all.dns_om2_phys(S)
            if S.backend == "gpu":
                import cupy as cp  # type: ignore
                field = cp.asnumpy(S.ur_full[2, :, :])
            else:
                field = np.asarray(S.ur_full[2, :, :])
            plane = self._float_to_pixels(field)
        elif var == self.VAR_STREAM:
            # Use dns_all stream-function helper: fills ur_full[2,:,:]
            dns_all.dns_stream_func(S)
            if S.backend == "gpu":
                import cupy as cp  # type: ignore
                field = cp.asnumpy(S.ur_full[2, :, :])
            else:
                field = np.asarray(S.ur_full[2, :, :])
            plane = self._float_to_pixels(field)
        else:
            plane = self._snapshot(1)

        return plane

    # ------------------------------------------------------------------
    def get_frame_pixels(self) -> np.ndarray:
        """Used by the Qt app worker thread.

        FIELD2PIX / dns_frame currently returns 32-bit packed gray pixels
        (0x00LLLLLL). Here we reduce that once to an 8-bit contiguous
        array so the GUI can push it straight into a QImage.
        """
        plane = self.make_pixels_component(self.current_var)

        # plane is already uint8 [0..255]; we keep the original logic
        #   pixels32 & 0xFF → pixels8
        pixels32 = np.asarray(plane, dtype=np.uint32)
        pixels8 = (pixels32 & 0xFF).astype(np.uint8)
        return np.ascontiguousarray(pixels8)

    def set_variable(self, var: int) -> None:
        """Select which variable the GUI should visualize."""
        self.current_var = int(var)

    def get_time(self) -> float:
        return float(self.t)

    def get_iteration(self) -> int:
        return self.iteration


    def diagnostics(self) -> dict:
        return {"t": float(self.t), "dt": float(self.dt), "cn": float(self.cn)}

def main():
    sim = DnsSimulator()
    print("Starting DNS simulation")
    for i in range(10):
        sim.step(1)
        d = sim.diagnostics()
        print(f"Step {i + 1:2d}: T={d['t']:.6f}, DT={d['dt']:.6f}, CN={d['cn']:.6f}")

    pix = sim.make_pixels_component()
    print("Pixels:", pix.shape, pix[0, 0], pix[10, 10])

# ----------------------------------------------------------------------
if __name__ == "__main__":
    main()