import glob
import json
import os
from typing import Any, Literal

import h5py
import numpy as np

from cqpes.types import TrainConfig
from cqpes.utils.train import find_best_checkpoint
from cqpes.utils.workspace import ExperimentWorkspace

_ACT_MAP = {
    "linear": 0,
    "tanh": 1,
    "sigmoid": 2,
    "isru": 3,
}


def _model2potfit(
    model: Any,  # tf.keras.Model
    output_dir: str,
    p_min: np.ndarray,
    p_max: np.ndarray,
    V_min: float,
    V_max: float,
) -> tuple[str, str]:
    weights_file = os.path.join(output_dir, "weights.txt")
    biases_file = os.path.join(output_dir, "biases.txt")

    # weights.txt
    with open(weights_file, "w") as f:
        print(
            f"{model.input_shape[1]:8d}"
            f"{(len(model.layers) - 1):8d}"
            f"{model.output_shape[1]:8d}",
            file=f,
        )

        for idx, layer in enumerate(model.layers):
            if idx == (len(model.layers) - 1):
                print("", file=f)

                break

            units = layer.get_config()["units"]

            print(f"{units:8d}", end="", file=f)

        act_name = model.layers[0].get_config()["activation"]
        act_id = _ACT_MAP.get(act_name)

        if act_id is None:
            raise ValueError(
                f"Activation '{act_name}' not supported for txt export."
            )

        print(f"{act_id:8d}{model.count_params():8d}", file=f)

        pdela = 0.5 * (
            np.concatenate((p_max[1:], [V_max]))
            - np.concatenate((p_min[1:], [V_min]))
        )

        pavga = 0.5 * (
            np.concatenate((p_max[1:], [V_max]))
            + np.concatenate((p_min[1:], [V_min]))
        )

        for x in pdela:
            print(f"{x:20.16f}\t", end="", file=f)

        print("", file=f)

        for x in pavga:
            print(f"{x:20.16f}\t", end="", file=f)

        print("", file=f)

        for layer in model.layers:
            W, _ = layer.get_weights()

            for w in W.T.flatten():
                print(f"{w:40.20f}", file=f)

    # bias.txt
    with open(biases_file, "w") as f:
        for layer in model.layers:
            _, b = layer.get_weights()

            for x in b:
                print(f"{x:40.20f}", file=f)

    return weights_file, biases_file


def _model2keras(
    model: Any,  # tf.keras.Model
    output_dir: str,
    alpha: float,
    ref_energy: float,
    p_min: np.ndarray,
    p_max: np.ndarray,
    V_min: float,
    V_max: float,
) -> str:
    h5_path = os.path.join(output_dir, "model.h5")

    # save keras weights
    model.save(h5_path)

    # save phys
    with h5py.File(h5_path, "a") as f:
        g = f.require_group("phys_metadata")

        g.attrs["alpha"] = alpha
        g.attrs["ref_energy"] = ref_energy

        g.attrs["V_min"] = V_min
        g.attrs["V_max"] = V_max

        for key, val in {"p_min": p_min, "p_max": p_max}.items():
            if key in g:
                del g[key]

            g.create_dataset(name=key, data=val)

    return h5_path


def _model2jaxpip(
    model: Any,  # tf.keras.Model
    basis_file: str,
    output_dir: str,
    alpha: float,
    p_min: np.ndarray,
    p_max: np.ndarray,
    V_min: float,
    V_max: float,
) -> str:
    from cqpes._env import _setup_jax

    _setup_jax()

    import equinox as eqx
    import jax
    from jax import numpy as jnp

    from jaxpip.descriptor import PolynomialDescriptor
    from jaxpip.model import FeatureScaler, PolynomialNeuralNetwork

    descriptor = PolynomialDescriptor.from_file(
        basis_file=basis_file,
        alpha=alpha,
        decay_kernel="morse",  # TODO: reciprocal
        dtype=jnp.float64,
    )

    hidden_units = [layer.units for layer in model.layers[:-1]]

    skeleton = PolynomialNeuralNetwork(
        descriptor=descriptor,
        hidden_layers=hidden_units,
        key=jax.random.PRNGKey(114514),
        activation=model.layers[0].get_config()["activation"],
    )

    new_model = skeleton

    linear_indices = [
        i
        for i, layer in enumerate(new_model.layers.layers)
        if isinstance(layer, eqx.nn.Linear)
    ]

    for k_idx, eqx_idx in enumerate(linear_indices):
        W_keras, b_keras = model.layers[k_idx].get_weights()

        # NOTE: Transpose required
        W_jax = jnp.array(W_keras.T, dtype=jnp.float64)
        b_jax = jnp.array(b_keras, dtype=jnp.float64)

        new_model = eqx.tree_at(
            where=lambda m, idx=eqx_idx: (
                m.layers.layers[idx].weight,
                m.layers.layers[idx].bias,
            ),
            pytree=new_model,
            replace=(W_jax, b_jax),
        )

    new_scaler = FeatureScaler(
        p_min=jnp.array(p_min[1:], dtype=jnp.float64),
        p_max=jnp.array(p_max[1:], dtype=jnp.float64),
        V_min=jnp.array(V_min, dtype=jnp.float64),
        V_max=jnp.array(V_max, dtype=jnp.float64),
    )

    new_model = eqx.tree_at(
        where=lambda m: m.scaler,
        pytree=new_model,
        replace=new_scaler,
    )

    eqx_path = os.path.join(output_dir, "jaxpip_network.eqx")
    new_model.save(eqx_path)

    return eqx_path


def run_export(
    workdir_path: str,
    export_type: Literal["h5", "potfit", "jaxpip"],
) -> None:
    # lazy import
    from cqpes._env import _setup_tensorflow

    _setup_tensorflow()

    import tf_levenberg_marquardt as lm

    from cqpes.utils.model import build_network

    # 1. existing workspace
    workspace = ExperimentWorkspace.from_existing(workdir_path)
    config_path = os.path.join(workspace.path, "train.json")

    # 2. load train config
    train_config = TrainConfig.from_json(config_path)

    # find jaxpip basis
    if export_type == "jaxpip":
        json_files = glob.glob(os.path.join(train_config.data, "*.json.gz"))

        if len(json_files) == 0:
            raise RuntimeError("Error: No JaxPIP basis json found")
        elif len(json_files) == 1:
            basis_file = json_files[0]
        elif len(json_files) > 1:
            raise RuntimeError(
                f"Error: Multiple JaxPIP basis json found: {json_files}"
            )

    # 3. load morse range parameter &  scale factor
    alpha = np.load(os.path.join(workspace.path, "alpha.npy")).item()
    ref_energy = np.load(os.path.join(workspace.path, "ref_energy.npy")).item()

    phys_dict = {
        k: np.load(os.path.join(workspace.path, f"{k}.npy"))
        for k in ["p_min", "p_max", "V_min", "V_max"]
    }

    for k in ["V_min", "V_max"]:
        phys_dict[k] = phys_dict[k].item()

    # 4. auto detect checkpoint
    best_ckpt_path, _ = find_best_checkpoint(workspace.path)

    # 5. build model
    input_dim = len(phys_dict["p_min"]) - 1
    model = build_network(train_config, input_dim=input_dim)
    model_wrapper = lm.model.ModelWrapper(model)  # type: ignore
    model_wrapper.build(input_shape=(1, input_dim))
    model_wrapper.load_weights(best_ckpt_path)

    # 6. export
    export_dir = workspace.get_subpath("export")
    os.makedirs(export_dir, exist_ok=True)

    if export_type == "h5":
        model_h5 = _model2keras(
            model,
            output_dir=export_dir,
            alpha=alpha,
            ref_energy=ref_energy,
            **phys_dict,
        )

        print(
            f"  [{'H5':^10}] Saved pure Keras model to: "
            f"{os.path.basename(model_h5)}"
        )
    elif export_type == "potfit":
        weights_path, biases_path = _model2potfit(
            model_wrapper.model,
            export_dir,
            **phys_dict,
        )

        print(
            f"  [{'POTFIT':^10}] Saved potfit model to "
            f"{os.path.basename(weights_path)} and "
            f"{os.path.basename(biases_path)}"
        )
    elif export_type == "jaxpip":
        model_eqx = _model2jaxpip(
            model,
            basis_file=basis_file,
            output_dir=export_dir,
            alpha=alpha,
            **phys_dict,
        )

        print(
            f"  [{'JAXPIP':^10}] Saved pure Equinox model to: "
            f"{os.path.basename(model_eqx)}"
        )

    meta_path = os.path.join(export_dir, "model_info.json")

    hidden_units = [layer.units for layer in model.layers[:-1]]

    with open(meta_path, "w") as f:
        json.dump(
            {
                "alpha": alpha,
                "ref_energy": ref_energy,
                "decay_kernel": "morse",  # TODO: reciprocal
                "feature_dim": model.input_shape[1],
                "hidden_layers": hidden_units,
                "activation": model.layers[0].get_config()["activation"],
            },
            f,
            indent=4,
        )

    print(f"  [{'DONE':^10}] Results saved to {export_dir}")
