from __future__ import annotations

import os
from typing import Optional, Sequence, Tuple

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
import warnings
import logging
import numpy as np
import cv2
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

class _EfficientNetBackbone(nn.Module):
    _LAST_CHANNELS: dict[str, int] = {
        'b0': 1280, 'b1': 1280, 'b2': 1408,
        'b3': 1536, 'b4': 1792, 'b5': 2048,
        'b6': 2304, 'b7': 2560,
    }

    def __init__(
        self,
        variant: str = 'b5',
        pretrained: bool = False,
        out_indices: Tuple[int, ...] = (8,),
        frozen_stages: int = -1,
        norm_eval: bool = False,
    ) -> None:
        super().__init__()

        variant = variant.lower()
        assert variant in self._LAST_CHANNELS, (
            f"Unknown EfficientNet variant '{variant}'. "
            f"Choose from {list(self._LAST_CHANNELS)}."
        )

        self.variant = variant
        self.out_indices = out_indices
        self.frozen_stages = frozen_stages
        self.norm_eval = norm_eval
        self.out_channels = self._LAST_CHANNELS[variant]
        import torchvision.models as tvm

        weights_arg = 'DEFAULT' if pretrained else None
        builder = getattr(tvm, f'efficientnet_{variant}')
        is_dist = dist.is_available() and dist.is_initialized()
        local_rank = int(os.environ.get('LOCAL_RANK', 0))
        if is_dist and local_rank != 0:
            dist.barrier()
        tv_model = builder(weights=weights_arg)
        if is_dist and local_rank == 0:
            dist.barrier()

        self.features: nn.Sequential = tv_model.features
        self.classifier = tv_model.classifier

        self._freeze_stages()

    def _freeze_stages(self) -> None:
        for i, layer in enumerate(self.features):
            if i <= self.frozen_stages:
                layer.eval()
                for param in layer.parameters():
                    param.requires_grad = False

    def train(self, mode: bool = True) -> 'EfficientNetBackbone':
        super().train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)):
                    m.eval()
        return self

    def forward(self, x: Tensor) -> Tuple[Tensor, ...]:
        outs = []
        for i, layer in enumerate(self.features):
            x = layer(x)
            if i in self.out_indices:
                outs.append(x)
        return tuple(outs)

class HeatmapHead(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        deconv_out_channels: Sequence[int] = (256, 256, 256),
        deconv_kernel_sizes: Sequence[int] = (4, 4, 4),
        conv_out_channels: Optional[Sequence[int]] = None,
        conv_kernel_sizes: Optional[Sequence[int]] = None,
        final_kernel_size: int = 1,
    ) -> None:
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels

        if deconv_out_channels:
            assert len(deconv_out_channels) == len(deconv_kernel_sizes), (
                "'deconv_out_channels' and 'deconv_kernel_sizes' must have "
                "equal length."
            )
            self.deconv_layers = self._make_deconv_layers(
                in_channels, deconv_out_channels, deconv_kernel_sizes
            )
            in_channels = deconv_out_channels[-1]
        else:
            self.deconv_layers = nn.Identity()

        if conv_out_channels:
            assert conv_kernel_sizes is not None and len(
                conv_out_channels) == len(conv_kernel_sizes), (
                "'conv_out_channels' and 'conv_kernel_sizes' must have "
                "equal length."
            )
            self.conv_layers = self._make_conv_layers(
                in_channels, conv_out_channels, conv_kernel_sizes
            )
            in_channels = conv_out_channels[-1]
        else:
            self.conv_layers = nn.Identity()

        pad = (final_kernel_size - 1) // 2
        self.final_layer = nn.Conv2d(
            in_channels, out_channels,
            kernel_size=final_kernel_size,
            padding=pad,
        )

        self._init_weights()

    def _init_weights(self) -> None:
        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                nn.init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    @staticmethod
    def _make_deconv_layers(
        in_channels: int,
        out_channels_list: Sequence[int],
        kernel_sizes: Sequence[int],
    ) -> nn.Sequential:
        layers: list[nn.Module] = []
        for out_ch, ks in zip(out_channels_list, kernel_sizes):
            if ks == 4:
                padding, output_padding = 1, 0
            elif ks == 3:
                padding, output_padding = 1, 1
            elif ks == 2:
                padding, output_padding = 0, 0
            else:
                raise ValueError(
                    f"Unsupported deconv kernel size {ks}. Use 2, 3, or 4."
                )
            layers += [
                nn.ConvTranspose2d(
                    in_channels, out_ch,
                    kernel_size=ks, stride=2,
                    padding=padding, output_padding=output_padding,
                    bias=False,
                ),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace=True),
            ]
            in_channels = out_ch
        return nn.Sequential(*layers)

    @staticmethod
    def _make_conv_layers(
        in_channels: int,
        out_channels_list: Sequence[int],
        kernel_sizes: Sequence[int],
    ) -> nn.Sequential:
        layers: list[nn.Module] = []
        for out_ch, ks in zip(out_channels_list, kernel_sizes):
            padding = (ks - 1) // 2
            layers += [
                nn.Conv2d(in_channels, out_ch,
                          kernel_size=ks, stride=1, padding=padding),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace=True),
            ]
            in_channels = out_ch
        return nn.Sequential(*layers)

    def forward(self, x: Tensor) -> Tensor:
        x = self.deconv_layers(x)
        x = self.conv_layers(x)
        x = self.final_layer(x)
        return x

class EfficientNetB5PoseNet(nn.Module):
    def __init__(
        self,
        num_keypoints: int = 17,
        pretrained: bool = False,
        frozen_stages: int = -1,
        norm_eval: bool = False,
        deconv_out_channels: Tuple[int, ...] = (256, 256, 256),
        deconv_kernel_sizes: Tuple[int, ...] = (4, 4, 4),
    ) -> None:
        super().__init__()

        self.backbone = _EfficientNetBackbone(
            variant='b5',
            pretrained=pretrained,
            out_indices=(8,),
            frozen_stages=frozen_stages,
            norm_eval=norm_eval,
        )
        backbone_out_ch = self.backbone.out_channels
        self.head = HeatmapHead(
            in_channels=backbone_out_ch,
            out_channels=num_keypoints,
            deconv_out_channels=deconv_out_channels,
            deconv_kernel_sizes=deconv_kernel_sizes,
        )

    def forward(self, x: Tensor) -> Tensor:
        feats: Tuple[Tensor, ...] = self.backbone(x)
        feat: Tensor = feats[-1]
        heatmaps: Tensor = self.head(feat)
        return heatmaps


DEFAULT_INPUT_SIZE = (192, 256)

class PoseEstimator:
    def __init__(self, model_name, num_keypoints=17, device=None, input_size=DEFAULT_INPUT_SIZE):
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device)
        self.input_size = input_size
        self.model_name = model_name
        self.model = EfficientNetB5PoseNet(num_keypoints=num_keypoints)
        if os.path.isfile(model_name):
            weights_path = model_name
        elif os.path.isdir(model_name):
            weights_path = os.path.join(model_name, "model.safetensors")
        else:
            weights_path = hf_hub_download(repo_id=model_name, filename="model.safetensors")
        state_dict = load_file(weights_path, device=str(self.device))
        self.model.load_state_dict(state_dict, strict=False)
        self.model.to(self.device)
        self.model.eval()
        self.num_keypoints = num_keypoints

    @staticmethod
    def _get_centers_and_scales_xyxy(person_boxes, scale_factor=1.0):
        centers, scales = [], []
        for box in person_boxes:
            x1, y1, x2, y2 = box
            x1, x2 = sorted([x1, x2])
            y1, y2 = sorted([y1, y2])
            centers.append([(x1+x2)/2.0, (y1+y2)/2.0])
            w, h = x2-x1, y2-y1
            scales.append([(w/200.0)*scale_factor, (h/200.0)*scale_factor])
        return np.array(centers), np.array(scales)

    @staticmethod
    def _process_image(image, bbox, target_size, angle=0, flip=False):
        try:
            if image is None or not isinstance(image, np.ndarray):
                raise ValueError("Invalid image input.")
            x1, y1, x2, y2 = map(lambda v: int(round(v)), bbox)
            if x2-x1 <= 0 or y2-y1 <= 0:
                raise ValueError(f"Invalid bbox: {{bbox}}")
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)
            if x2 <= x1 or y2 <= y1:
                raise ValueError("Invalid bbox after clamping.")
            cropped = image[y1:y2, x1:x2]
            resized = cv2.resize(cropped, target_size)
            if angle != 0:
                center = (target_size[0]//2, target_size[1]//2)
                rot = cv2.getRotationMatrix2D(center, angle, 1.0)
                resized = cv2.warpAffine(resized, rot, target_size)
            if flip:
                resized = cv2.flip(resized, 1)
            return resized, True
        except Exception:
            blank = np.zeros((target_size[1], target_size[0], 3), dtype=np.uint8)
            return blank, False

    @staticmethod
    def _process(image, target_size=(192, 256), angle=0, flip=False, conf_threshold=0.5, model_weights="yolov8n.pt"):
        try:
            from ultralytics import YOLO
        except ImportError:
            raise ImportError("ultralytics is required. pip install ultralytics")
        model = YOLO(model_weights)
        crops, metadata = [], []
        if image is None or not isinstance(image, np.ndarray):
            raise ValueError("Invalid image input.")
        results = model(image, conf=conf_threshold, classes=[0], verbose=False)
        bboxes = []
        for r in results:
            for box in r.boxes:
                bboxes.append(box.xyxy[0].cpu().numpy().tolist())
        for idx, bbox in enumerate(bboxes):
            processed, success = PoseEstimator._process_image(image, bbox, target_size, angle, flip)
            crops.append(processed)
            metadata.append({"bbox": bbox, "person_index": idx, "success": success})
        if not crops:
            return None, metadata
        batch = np.stack(crops, axis=0).transpose(0, 3, 1, 2)
        return np.ascontiguousarray(batch), metadata

    def _preprocess(self, image_bgr):
        batch, meta = self._process(image_bgr, target_size=self.input_size)
        if batch is None:
            return None, meta
        t = torch.tensor(batch, dtype=torch.float32) / 255.0
        return t.to(self.device), meta

    @staticmethod
    def _taylor(heatmap, coord):
        H, W = heatmap.shape[:2]
        px, py = int(coord[0]), int(coord[1])
        if 1 < px < W-2 and 1 < py < H-2:
            dx = 0.5*(heatmap[py][px+1]-heatmap[py][px-1])
            dy = 0.5*(heatmap[py+1][px]-heatmap[py-1][px])
            dxx = 0.25*(heatmap[py][px+2]-2*heatmap[py][px]+heatmap[py][px-2])
            dxy = 0.25*(heatmap[py+1][px+1]-heatmap[py-1][px+1]-heatmap[py+1][px-1]+heatmap[py-1][px-1])
            dyy = 0.25*(heatmap[py+2][px]-2*heatmap[py][px]+heatmap[py-2][px])
            derivative = np.array([[dx],[dy]])
            hessian = np.array([[dxx,dxy],[dxy,dyy]])
            if dxx*dyy - dxy**2 != 0:
                offset = -np.linalg.inv(hessian) @ derivative
                coord += np.squeeze(offset.T, axis=0)
        return coord

    @staticmethod
    def _get_max_preds(heatmaps):
        N, K, _, W = heatmaps.shape
        reshaped = heatmaps.reshape((N, K, -1))
        idx = np.argmax(reshaped, 2).reshape((N, K, 1))
        maxvals = np.amax(reshaped, 2).reshape((N, K, 1))
        preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
        preds[:,:,0] = preds[:,:,0] % W
        preds[:,:,1] = np.floor(preds[:,:,1] / W)
        preds = np.where(np.tile(maxvals, (1, 1, 2)) > 0.0, preds, -1)
        return preds, maxvals

    @staticmethod
    def _gaussian_blur(heatmaps, kernel=11):
        border = (kernel-1)//2
        B, J, H, W = heatmaps.shape
        for i in range(B):
            for j in range(J):
                origin_max = np.max(heatmaps[i,j])
                dr = np.zeros((H+2*border, W+2*border), dtype=np.float32)
                dr[border:-border, border:-border] = heatmaps[i,j].copy()
                dr = cv2.GaussianBlur(dr, (kernel, kernel), 0)
                heatmaps[i,j] = dr[border:-border, border:-border].copy()
                heatmaps[i,j] *= origin_max / np.max(heatmaps[i,j])
        return heatmaps

    @staticmethod
    def transform_preds(coords, center, scale, output_size, use_udp=False):
        scale = scale * 200.0
        if use_udp:
            sx = scale[0]/(output_size[0]-1.0)
            sy = scale[1]/(output_size[1]-1.0)
        else:
            sx = scale[0]/output_size[0]
            sy = scale[1]/output_size[1]
        tc = np.ones_like(coords)
        tc[:,0] = coords[:,0]*sx + center[0] - scale[0]*0.5
        tc[:,1] = coords[:,1]*sy + center[1] - scale[1]*0.5
        return tc

    @staticmethod
    def keypoints_from_heatmaps(heatmaps, center, scale, unbiased=False, post_process="default", kernel=11, use_udp=False, target_type="GaussianHeatmap"):
        heatmaps = heatmaps.copy()
        if unbiased:
            assert post_process not in [False, None, "megvii"]
        if post_process == "default" and unbiased:
            post_process = "unbiased"
        if post_process == "megvii":
            heatmaps = PoseEstimator._gaussian_blur(heatmaps, kernel=kernel)
        N, K, H, W = heatmaps.shape
        preds, maxvals = PoseEstimator._get_max_preds(heatmaps)
        if post_process == "unbiased":
            heatmaps = np.log(np.maximum(PoseEstimator._gaussian_blur(heatmaps, kernel), 1e-10))
            for n in range(N):
                for k in range(K):
                    preds[n][k] = PoseEstimator._taylor(heatmaps[n][k], preds[n][k])
        elif post_process is not None and post_process != "megvii":
            for n in range(N):
                for k in range(K):
                    hm = heatmaps[n][k]
                    px, py = int(preds[n][k][0]), int(preds[n][k][1])
                    if 1 < px < W-1 and 1 < py < H-1:
                        diff = np.array([hm[py][px+1]-hm[py][px-1], hm[py+1][px]-hm[py-1][px]])
                        preds[n][k] += np.sign(diff)*0.25
        for i in range(N):
            preds[i] = PoseEstimator.transform_preds(preds[i], center[i], scale[i], [W, H], use_udp=use_udp)
        if post_process == "megvii":
            maxvals = maxvals/255.0 + 0.5
        return preds, maxvals

    @torch.no_grad()
    def predict(self, image_bgr):
        tensor, meta = self._preprocess(image_bgr)
        if tensor is None:
            return np.array([]), np.array([])
        centers, scales = self._get_centers_and_scales_xyxy([m["bbox"] for m in meta])
        output = self.model(tensor).detach().cpu().numpy()
        kps, scores = self.keypoints_from_heatmaps(output, centers, scales, unbiased=True, post_process="default", target_type="GaussianHeatmap", kernel=11)
        return kps, scores

    @staticmethod
    def visualize(image_bgr, keypoints, scores, score_threshold=0.3, kp_radius=8, line_thickness=5):
        canvas = image_bgr.copy()
        if keypoints.ndim == 2:
            keypoints = np.expand_dims(keypoints, axis=0)
            scores = np.expand_dims(scores, axis=0)
        edges = [(0,1),(0,2),(1,3),(2,4),(5,6),(5,11),(6,12),(11,12),(5,7),(7,9),(6,8),(8,10),(11,13),(13,15),(12,14),(14,16)]
        colors = [(255,0,0),(255,85,0),(255,170,0),(255,255,0),(170,255,0),(85,255,0),(0,255,0),(0,255,85),(0,255,170),(0,255,255),(0,170,255),(0,85,255),(0,0,255),(85,0,255),(170,0,255),(255,0,255)]
        for n in range(len(keypoints)):
            kpts, scs = keypoints[n], scores[n].squeeze()
            for i,(a,b) in enumerate(edges):
                if scs[a]>=score_threshold and scs[b]>=score_threshold:
                    cv2.line(canvas,(int(kpts[a][0]),int(kpts[a][1])),(int(kpts[b][0]),int(kpts[b][1])),colors[i],thickness=line_thickness)
            for k in range(len(kpts)):
                if scs[k]>=score_threshold:
                    cv2.circle(canvas,(int(kpts[k,0]),int(kpts[k,1])),kp_radius,color=(255,255,255),thickness=-1)
        return canvas