from typing import Callable, Dict, List, Literal, Optional

# import matplotlib.pyplot as plt
import torch
from torchmetrics import Metric, MetricCollection

from zyplib.train.utils import use_device


class TorchMetricRecorder:
    def __init__(
        self,
        metrics: MetricCollection | dict[str, Metric],
        before_metric: Literal['sigmoid', 'softmax']
        | Callable[[torch.Tensor], torch.Tensor]
        | dict[str | Literal['*'], Callable] = None,
    ):
        """添加基于 `torchmetrics` 的日志评估指标的功能

        Parameters
        ----------
        - `metrics` : `MetricCollection | Dict[str, Metric]`, optional
            - 评估指标
            - 使用 `torchmetrics.MetricCollection` 或 `dict` 类型
        - `before_metric` : `'sigmoid' | 'softmax' | Callable | dict[str, Callable]`, optional

        `before_metric`:
        ----------
        在计算评估指标之前的激活函数: `(y_hat) -> y_activated`

        例如在二分类问题中，模型最后一层没有激活函数, 在计算评估指标 (即 `metrics`) 之前需要使用 `torch.sigmoid` 激活

        - `None`: 不需要额外的激活函数
        - `Literal['sigmoid', 'softmax']`: 使用 torch.sigmoid 和 torch.softmax
        - `(Tensor) -> Tensor`: 指定的激活函数
        - `dict[str, (Tensor) -> Tensor]`: 在比较复杂的情况下, 各个指标使用的输入不同的情况下可以用这个
            - key: 如果为 `*` 则代表通配符
            - value: `(Tensor) -> Tensor`, 为 None，则不激活
        """
        if before_metric == 'sigmoid':
            self.before_metric = torch.sigmoid
        elif before_metric == 'softmax':
            self.before_metric = torch.softmax
        else:
            self.before_metric = before_metric

        if metrics is None:
            self.metrics = MetricCollection({})
        elif isinstance(metrics, dict):
            self.metrics = MetricCollection(metrics)
        else:
            self.metrics = metrics

        self.step_history = {metric_name: [] for metric_name in self.metrics.keys()}
        self.step_history['loss'] = []
        self.epoch_history = []

        self.current_epoch = 0
        self.current_step = 0

        self.epoch_metrics = {
            metric_name: metric.clone() for metric_name, metric in self.metrics.items()
        }
        self.epoch_losses = []

    def to_device(self, device: str | torch.device):
        for metrics in [self.metrics, self.epoch_metrics]:
            for key, metric in metrics.items():
                # self.metrics[key] = use_device(metric, device)
                use_device(metric, device)

    def update(self, y_hat, y, loss: Optional[float] = None):
        """更新某个 step 的评估指标

        Parameters
        ----------
        - `y_hat` : `torch.Tensor`
            - 预测值
        - `y` : `torch.Tensor`
            - 真实值
        - `loss` : `float?`, optional
            - 损失值, 注意是 `float` 类型, 需要使用 `item()` 方法获取
        """
        result = {}

        # Log loss if provided
        if loss is not None:
            if isinstance(loss, torch.Tensor):
                loss = loss.item()
            self.step_history['loss'].append(loss)
            self.epoch_losses.append(loss)
            result['loss'] = loss

        # 处理 before_metric 激活
        if self.before_metric is None:
            y_hat_activated = y_hat
        elif isinstance(self.before_metric, dict):
            # 如果是字典类型，需要对不同的指标使用不同的激活函数
            y_hat_activated = {}
            for metric_name in self.metrics.keys():
                # 检查是否有针对特定指标的激活函数
                if metric_name in self.before_metric:
                    activation_fn = self.before_metric[metric_name]
                # 检查是否有通配符激活函数
                elif '*' in self.before_metric:
                    activation_fn = self.before_metric['*']
                else:
                    activation_fn = None

                if activation_fn is not None:
                    y_hat_activated[metric_name] = activation_fn(y_hat)
                else:
                    y_hat_activated[metric_name] = y_hat
        else:
            # 如果是单个激活函数
            if self.before_metric == torch.softmax:
                # softmax 需要指定 dim 参数
                y_hat_activated = self.before_metric(y_hat, dim=-1)
            else:
                y_hat_activated = self.before_metric(y_hat)

        # Compute and log metrics
        for metric_name, metric in self.metrics.items():
            # 如果 y_hat_activated 是字典，使用对应的激活结果
            if isinstance(y_hat_activated, dict):
                current_y_hat = y_hat_activated[metric_name]
            else:
                current_y_hat = y_hat_activated

            step_metric_value = metric(current_y_hat, y)
            result[metric_name] = step_metric_value.item()
            self.step_history[metric_name].append(step_metric_value.item())
            # Update epoch-level metrics
            self.epoch_metrics[metric_name].update(current_y_hat, y)

        self.current_step += 1
        return result

    def compute(self):
        """计算一个 epoch 的评估指标; 在每个 epoch 结束时调用"""
        epoch_metrics = {}

        # Compute epoch-level loss
        if self.epoch_losses:
            epoch_loss = sum(self.epoch_losses) / len(self.epoch_losses)
            epoch_metrics['loss'] = epoch_loss
            # self.fn_logging('epoch_loss', epoch_loss)

        # Compute and log epoch-level metrics
        for metric_name, metric in self.epoch_metrics.items():
            epoch_metric_value = metric.compute()
            epoch_metrics[metric_name] = epoch_metric_value.item()
            # self.fn_logging(f'epoch_{metric_name}', epoch_metric_value.item())
            metric.reset()

        self.epoch_history.append(epoch_metrics)
        self.current_epoch += 1
        self.current_step = 0
        self.epoch_losses = []
        return epoch_metrics

    def get_history(
        self, metric_name: str, level: Literal['step', 'epoch'] = 'step'
    ) -> List[float]:
        if level == 'step':
            return self.step_history[metric_name]
        elif level == 'epoch':
            return [
                epoch[metric_name] for epoch in self.epoch_history if metric_name in epoch
            ]
        else:
            raise ValueError("Level must be either 'step' or 'epoch'")
