from typing import List

import torch
from torch import Tensor
from torchmetrics import Metric

from .utils import *


# motion reconstruction metric
class PredMetrics(Metric):

    def __init__(self,
                 cfg,
                 njoints: int = 22,
                 jointstype: str = "mmm",
                 force_in_meter: bool = True,
                 align_root: bool = True,
                 dist_sync_on_step=True,
                 task: str = "pred",
                 **kwargs):
        super().__init__(dist_sync_on_step=dist_sync_on_step)

        self.name = 'Motion Prdiction'
        self.cfg = cfg
        self.jointstype = jointstype
        self.align_root = align_root
        self.task = task
        self.force_in_meter = force_in_meter

        self.add_state("count", default=torch.tensor(0), dist_reduce_fx="sum")
        self.add_state("count_seq",
                       default=torch.tensor(0),
                       dist_reduce_fx="sum")

        self.add_state("APD",
                       default=torch.tensor([0.0]),
                       dist_reduce_fx="sum")
        self.add_state("ADE",
                       default=torch.tensor([0.0]),
                       dist_reduce_fx="sum")
        self.add_state("FDE",
                       default=torch.tensor([0.0]),
                       dist_reduce_fx="sum")

        self.MR_metrics = ["APD", "ADE", "FDE"]

        # All metric
        self.metrics = self.MR_metrics

    def compute(self, sanity_flag):

        count = self.count
        count_seq = self.count_seq
        mr_metrics = {}
        mr_metrics["APD"] = self.APD / count_seq 
        mr_metrics["ADE"] = self.ADE / count_seq
        mr_metrics["FDE"] = self.FDE / count_seq
        
        # Reset
        self.reset()
        
        return mr_metrics

    def update(self, joints_rst: Tensor, joints_ref: Tensor,
               lengths: List[int]):
        
        assert joints_rst.shape == joints_ref.shape
        assert joints_rst.dim() == 4
        # (bs, seq, njoint=22, 3)

        self.count += sum(lengths)
        self.count_seq += len(lengths)

        rst = torch.flatten(joints_rst, start_dim=2)
        ref = torch.flatten(joints_ref, start_dim=2)
        
        for i, l in enumerate(lengths):
            if self.task == "pred":
                pred_start = int(l*self.cfg.ABLATION.predict_ratio)
                diff = rst[i,pred_start:] - ref[i,pred_start:]
            elif self.task == "inbetween":
                inbetween_start = int(l*self.cfg.ABLATION.inbetween_ratio)
                inbetween_end = l - int(l*self.cfg.ABLATION.inbetween_ratio)
                diff = rst[i,inbetween_start:inbetween_end] - ref[i,inbetween_start:inbetween_end]
            else:
                print(f"Task {self.task} not implemented.")
                diff = rst - ref
            
            dist = torch.linalg.norm(diff, dim=-1)[None]

            ade = dist.mean(dim=1)
            fde = dist[:,-1]
            self.ADE = self.ADE + ade
            self.FDE = self.FDE + fde