import hashlib
import logging
import math
from pathlib import Path
from time import perf_counter
from typing import List, Literal, Optional, Sequence, Union

import numpy as np
import torch
import torch.utils.data
from torch import Tensor

import franken.metrics
from franken.metrics.base import BaseMetric
import franken.utils.distributed as dist_utils
from franken.data.base import Target
from franken.rf.model import FrankenPotential
from franken.trainers import BaseTrainer
from franken.trainers.log_utils import (
    DataSplit,
    HyperParameterGroup,
    LogCollection,
    LogEntry,
)
from franken.utils.linalg.cov import (
    lowmem_normalize_leading_eig,
    rank1_update,
    rankk_update,
)
from franken.utils.linalg.psdsolve import psd_ridge
from franken.utils.linalg.tri import triangular_lerp
from franken.utils.misc import params_grid, throughput


logger = logging.getLogger("franken")


class RandomFeaturesTrainer(BaseTrainer):
    def __init__(
        self,
        train_dataloader: torch.utils.data.DataLoader,
        random_features_normalization: Optional[Literal["leading_eig"]] = "leading_eig",
        log_dir: Path | None = None,  # If None, logging is disabled
        save_every_model: bool = True,
        device: Union[torch.device, str, int] = "cuda:0",
        dtype: Union[str, torch.dtype] = torch.float32,
        save_fmaps: bool = True,
    ):
        super().__init__(
            train_dataloader,
            log_dir=log_dir,
            save_every_model=save_every_model,
            device=device,
            dtype=dtype,
        )
        self.random_features_normalization = random_features_normalization
        self.save_fmaps = save_fmaps

    def on_fit_start(self, model: FrankenPotential):
        # initialize input scaler based on statistics property
        model.input_scaler.set_from_statistics(self.get_statistics(model)[0])
        # initialize energy shift based on atomic energies
        if not model.energy_shift.is_initialized:
            model.energy_shift.set_from_atomic_energies(
                self.train_dataloader.dataset.atomic_energies
            )

    def fit(
        self,
        model: FrankenPotential,
        solver_params: dict,
    ) -> tuple[LogCollection, torch.Tensor]:
        model = model.to(self.device)
        self.on_fit_start(model)
        model_hash = hashlib.md5(str(model.hyperparameters).encode())
        model_hash = model_hash.hexdigest()

        t_cov_coeffs_start = perf_counter()
        self._compute_covs_and_coeffs(model, self.train_dataloader)
        t_cov_coeffs = perf_counter() - t_cov_coeffs_start

        solver_grid_size = math.prod([len(v) for v in solver_params.values()])
        all_weights = torch.zeros(
            solver_grid_size,
            model.rf.total_random_features,
            dtype=self.buffer_dt,
            device=self.device,
        )

        solver_iter = throughput(
            params_grid(solver_params, split_distributed=True),
            desc="least-squares",
            units="models",
            device=self.device,
        )

        local_logs = dict()
        num_failed = torch.zeros((1,), device=self.device, dtype=torch.int)
        for hp_idx, hp_val in solver_iter:
            t_solve_start = perf_counter()
            try:
                weights = self._solve(**hp_val)
            except torch.linalg.LinAlgError as e:
                weights = torch.full_like(all_weights[hp_idx], torch.inf)
                num_failed += 1
                logger.debug(f"Hyperparameter {hp_val} failed. Error: {e}")
            finally:
                t_solve = perf_counter() - t_solve_start
            # Update all weights
            all_weights[hp_idx].copy_(weights.view(-1))

            # Logging
            solver_hps = hp_val | {"dtype": self.buffer_dt}
            hp_groups = model.hyperparameters | {"solver": solver_hps}
            hyperparameters = []
            for group_name, hps in hp_groups.items():
                hyperparameters.append(HyperParameterGroup.from_dict(group_name, hps))

            local_logs[hp_idx] = LogEntry(
                checkpoint_hash=model_hash,
                checkpoint_rf_weight_id=hp_idx,
                timings_cov_coeffs=t_cov_coeffs,
                timings_solve=t_solve,
                hyperparameters=hyperparameters,
            )

        # Broadcast weights and logs
        dist_utils.all_sum(all_weights)
        log_collection = LogCollection.gather_from_ranks(local_logs)
        assert len(log_collection) == solver_grid_size

        # Reporting failed runs
        dist_utils.all_sum(num_failed)

        if num_failed.item() > 0:
            logger.warning(
                f"Solver failed in {num_failed.item()}/{solver_grid_size} cases."
            )

        return log_collection, all_weights

    def evaluate(
        self,
        model: FrankenPotential,
        dataloader: torch.utils.data.DataLoader,
        log_collection: LogCollection,
        all_weights: torch.Tensor | None,
        metrics: Sequence[str] = ("energy_MAE", "forces_MAE", "forces_cosim"),
    ) -> LogCollection:
        tot_dset_size = len(dataloader.dataset)  # type: ignore

        metric_objects: List[BaseMetric] = []
        for name in metrics:
            try:
                metric_obj = franken.metrics.init_metric(
                    name, device=self.device, dtype=self.buffer_dt
                )
                metric_objects.append(metric_obj)
            except KeyError:
                logger.warning(
                    f"Unknown metric {name}. Skipping. Available metrics: {franken.metrics.available_metrics()}"
                )

        split_name = dataloader.dataset.split
        try:
            split = DataSplit[split_name.upper()]
        except KeyError:
            logger.warning(f"Unrecognized split '{split_name}' in dataloader")
            split = DataSplit.UNDEFINED

        progress_bar = throughput(
            dataloader,
            desc=f"{split.name.lower()} evaluation",
            total=tot_dset_size,
            device=self.device,
        )
        with torch.no_grad():
            for i, (data, targets) in enumerate(progress_bar):
                data = data.to(device=self.device)
                targets = targets.to(device=self.device)
                if split == DataSplit.TRAIN and self.save_fmaps:
                    # Shortcut to compute predictions for the training-set, for which
                    # we already have computed energy and force feature maps
                    assert len(self.forces_fmap) == len(self.energy_fmap)
                    assert len(self.forces_fmap) == len(dataloader)
                    predictions = Target(
                        *model.energy_and_forces_from_fmaps(
                            data,
                            energy_fmap=self.energy_fmap[i],
                            forces_fmap=self.forces_fmap[i],
                            weights=all_weights,
                            add_energy_shift=False,  # since it's always train here
                        )
                    )
                else:
                    predictions = Target(
                        *model.energy_and_forces(
                            data,
                            weights=all_weights,
                            forces_mode="torch.func",
                            add_energy_shift=(
                                False if split == DataSplit.TRAIN else True
                            ),
                        )
                    )
                for metric in metric_objects:
                    metric.update(predictions, targets)

        num_models = (
            all_weights.shape[0]
            if all_weights is not None
            else model.rf.weights.shape[0]
        )
        # Sync metrics across GPUs
        metric_values = {}
        for metric in metric_objects:
            value = metric.compute(reset=False)
            if metric.samples_counter.item() != tot_dset_size:
                logger.warning(
                    f"Metric {metric.name} has {metric.samples_counter.item()} samples "
                    f"while the dataset has {tot_dset_size} configurations."
                )
            assert value.ndim == 1
            assert value.shape[0] == num_models
            metric_values[metric.name] = value

        # Explode metric values into list of MetricLog
        raw_logs = []
        for idx in range(num_models):
            results = []
            for name, value in metric_values.items():
                float_value = value[idx].item()
                results.append(dict(name=name, value=float_value))
            raw_logs.append(results)

        assert len(raw_logs) == len(log_collection)

        for log_entry, results in zip(log_collection, raw_logs):
            for metric in results:
                log_entry.add_metric(metric["name"], metric["value"], split)
        return log_collection

    def warn_save_fmaps(
        self, energy_fmap: Tensor, forces_fmap: Tensor, num_maps: int
    ) -> None:
        fmap_size = np.prod(energy_fmap.shape) + np.prod(forces_fmap.shape)
        tot_bytes = fmap_size * forces_fmap.element_size() * num_maps
        if self.device.type == "cuda":
            avail_bytes = torch.cuda.mem_get_info(self.device)[0]
            if tot_bytes > 0.8 * avail_bytes:
                logger.warning(
                    f"Saved feature maps require {tot_bytes / 2**30:.2f}GB of device memory. "
                    f"Device has {avail_bytes / 2**30:.2f}GB of available memory, a crash may occur. "
                    f"Use 'trainer.save_fmaps=False' in your config to stop saving feature maps."
                )
            else:
                logger.info(
                    f"Saved feature maps require {tot_bytes / 2**30:.2f}GB of device memory."
                )

    @torch.no_grad()
    def _compute_covs_and_coeffs(
        self, model: FrankenPotential, dataloader: torch.utils.data.DataLoader
    ) -> None:
        tot_dset_size = len(dataloader.dataset)  # type: ignore
        n_rf = model.rf.total_random_features
        # Initialize Buffers: `covariance` will have
        # - covariance of forces on LOWER triangle
        # - covariance of energies on UPPER triangle
        # with diagonals stored separately, and linsys coefficients.
        self.covariance = torch.zeros(
            (n_rf, n_rf), device=self.device, dtype=self.buffer_dt
        )
        self.diag_energy = torch.zeros(
            (n_rf,), device=self.device, dtype=self.buffer_dt
        )
        self.diag_forces = torch.zeros(
            (n_rf,), device=self.device, dtype=self.buffer_dt
        )
        self.coeffs_energy = torch.zeros(
            (n_rf,), device=self.device, dtype=self.buffer_dt
        )
        self.coeffs_forces = torch.zeros(
            (n_rf,), device=self.device, dtype=self.buffer_dt
        )
        # NOTE: the feature maps are never synced between devices.
        #       They can only used correctly by iterating through the
        #       same dataloader as here, using the same method. Otherwise
        #       they may not be in the correct order
        self.energy_fmap = []
        self.forces_fmap = []

        progress_bar = throughput(
            dataloader,
            desc="covs+coeffs",
            total=tot_dset_size,
            device=self.device,
        )

        for i, (data, targets) in enumerate(progress_bar):
            data = data.to(device=self.device)
            targets = targets.to(device=self.device)

            energy_per_atom = targets.energy / data.natoms
            forces_per_atom = targets.forces / data.natoms

            forces_fmap, energy_fmap = model.grad_feature_map(data)
            energy_fmap = energy_fmap.to(dtype=self.buffer_dt)

            rank1_update(self.covariance, self.diag_energy, energy_fmap, upper=True)
            self.coeffs_energy.add_(energy_fmap, alpha=energy_per_atom.item())

            forces_fmap = (-forces_fmap).to(dtype=self.buffer_dt)
            forces_fmap = forces_fmap.view(forces_fmap.shape[0], -1)
            rankk_update(self.covariance, self.diag_forces, forces_fmap, upper=False)
            self.coeffs_forces.addmv_(
                forces_fmap, forces_per_atom.view(-1).to(dtype=self.buffer_dt)
            )
            if self.save_fmaps:
                if i == 0:
                    self.warn_save_fmaps(energy_fmap, forces_fmap, len(dataloader))
                self.energy_fmap.append(energy_fmap)
                self.forces_fmap.append(forces_fmap)

        # Sync covariance matrices & coefficients
        dist_utils.all_sum(self.covariance)
        dist_utils.all_sum(self.diag_forces)
        dist_utils.all_sum(self.diag_energy)
        dist_utils.all_sum(self.coeffs_energy)
        dist_utils.all_sum(self.coeffs_forces)

        if self.random_features_normalization == "leading_eig":
            logger.warning(
                "`leading_eig` normalization has high memory usage. If you encounter OOM errors try to disable it."
            )
            lowmem_normalize_leading_eig(
                self.covariance, self.diag_energy, self.coeffs_energy, upper=True
            )
            lowmem_normalize_leading_eig(
                self.covariance, self.diag_forces, self.coeffs_forces, upper=False
            )
        elif self.random_features_normalization is not None:
            raise ValueError(
                f"Covariance normalization {self.random_features_normalization} is not implemented."
            )

    @torch.no_grad()
    def _solve(self, loss_lerp_weight: float, L2_penalty: float = 1e-6) -> Tensor:
        # This is the 2nd copy of the covariance matrix that we need to store.
        lerped_cov, lerped_diag = triangular_lerp(
            self.covariance,
            diag_upper=self.diag_energy,
            diag_lower=self.diag_forces,
            weight=loss_lerp_weight,
            inplace=False,
        )
        lerped_cov.diagonal().copy_(lerped_diag)
        rhs = torch.lerp(self.coeffs_energy, self.coeffs_forces, loss_lerp_weight)
        solution = psd_ridge(lerped_cov, rhs, L2_penalty)
        return solution
