from typing import Optional
from scipy.interpolate import interp1d

import datetime
import numpy as np
import torch
import kintera
import snapy


def write_profile(
    filename: str,
    hydro_w: torch.Tensor,
    block: snapy.MeshBlock,
    ref_pressure: float = 1.0e5,
    comment: Optional[str] = None,
) -> None:
    """
    Write an atmospheric profile to a text file.

    Args:
        filename (str): The name of the text file to write.
        hydro_w (torch.Tensor): Primitive variables of the atmosphere profile.
        block (snapy.MeshBlock): The mesh block containing the atmosphere data.
        ref_pressure (float): Reference pressure for height calculation.
        comment (Optional[str]): A comment string to write to the file.

    Returns:
        None
    """
    # useful modules
    thermo_y = block.module("hydro.eos.thermo")
    coord = block.module("hydro.coord")
    eos = block.hydro.get_eos()

    # handling mole fraction quantities
    thermo_x = kintera.ThermoX(thermo_y.options)
    thermo_x.to(dtype=hydro_w.dtype, device=hydro_w.device)

    # check dimensions of hydro_w
    species = thermo_y.options.species()
    if hydro_w.ndim != 4:
        raise ValueError("hydro_w must be a 4D tensor.")
    if hydro_w.shape[0] != 4 + len(species):
        raise ValueError("hydro_w shape does not match number of species.")
    if hydro_w.shape[1] != 1 or hydro_w.shape[2] != 1:
        raise ValueError("hydro_w must have shape (N, 1, 1, L).")

    # calculate a height grid
    pres = hydro_w[snapy.index.ipr, ...].squeeze() / 1.0e5  # Pa -> bar
    zlev_func = interp1d(
        pres.log().cpu().numpy(),
        coord.buffer("x1v").cpu().numpy(),
        kind="linear",
        fill_value="extrapolate",
    )
    zref = zlev_func(np.log(ref_pressure / 1.0e5))
    zlev = (coord.buffer("x1v") - zref) / 1.0e3  # m -> km

    # calculate temperature
    temp = eos.compute("W->T", (hydro_w,)).squeeze()

    # calculate mole fractions
    xfrac = thermo_y.compute("Y->X", (hydro_w[snapy.index.icy :, ...],)).squeeze()

    # calculate heat capacity
    conc = thermo_x.compute("TPX->V", (temp, pres * 1.0e5, xfrac))
    cpx = thermo_x.compute("TV->cp", (temp, conc)) / conc.sum(-1)

    # calculate entropy
    ens = thermo_x.compute("TPV->S", (temp, pres * 1.0e5, conc)) / conc.sum(-1)

    with open(filename, "w") as f:
        # write comments
        f.write(f"# File generated by write_profile.py\n")
        f.write(f"# Date generated: {datetime.datetime.now()}\n")
        f.write(f"# Speices units: mole fraction [ppmv]\n")
        f.write(f"# Width of the first column: 6 characters\n")
        f.write(f"# Width of other columns: 13 characters\n")
        f.write(f"# Reference pressure for z = 0: {ref_pressure} Pa\n")
        if comment:
            f.write(f"# {comment}\n\n")

        # Write header
        headers = []
        headers.append(f"# {'IDX':<4}")
        headers.append(f"{'HGT[km]':>{13}}")
        headers.append(f"{'PRE[bar]':>{13}}")
        headers.append(f"{'TEM[K]':>{13}}")
        for sp in species[1:]:
            headers.append(f"{sp.upper():>{13}}")
        headers.append(f"{'CPX[J,K,mol]':>{13}}")
        headers.append(f"{'ENS[J,K,mol]':>{13}}")

        f.write("".join(headers) + "\n")

        # Write data
        num_rows = hydro_w.shape[-1]
        for i in range(num_rows, 0, -1):
            row = []
            row.append(f"{i:<6}")
            row.append(f"{zlev[i-1]:.6g}".rjust(13))  # HGT
            row.append(f"{pres[i-1]:.6g}".rjust(13))  # PRE
            row.append(f"{temp[i-1]:.6g}".rjust(13))
            for s in range(1, len(species)):
                row.append(f"{xfrac[i-1,s]*1.e6:.6g}".rjust(13))
            row.append(f"{cpx[i-1]:.6g}".rjust(13))
            row.append(f"{ens[i-1]:.6g}".rjust(13))

            f.write("".join(row) + "\n")
    print(f"Atmosphere profile written to {filename}")
