import os, time
from typing import Dict, List, Optional, Union, Tuple, Any
from dataclasses import dataclass

import numpy as np

from .pipeline import (
    AutoBackend, 
    build_preprocess_pipeline, 
    build_postprocess_pipeline, 
    load_config,
    init_yolo_config,
    get_imgs_shapes,
    init_model_info
)



@dataclass(frozen=True)
class YOLOOptions:
    weights: str = './weights/yolo11s.pt'
    version: str = 'yolo11'
    task: str = 'detect'
    size: int = 320
    device: str = 'cpu'
    model_heads: Optional[List[int]] = None
    conf: float = 0.45
    iou: float = 0.5
    nc: int = 80
    agnostic: bool = False
    classes: Optional[List[int]] = None
    class_names: Optional[Dict[int, str]] = None
    fp16: bool = False
    kpt_shape: Optional[List[int]] = None
    pipeline: str = 'preprocess | inference | postprocess'
    streaming_mode: bool = False
    orig_imgs_shapes: Optional[List[Tuple[int, int]]] = None



class YOLO:
    def __init__(self, config: Union[str, YOLOOptions], streaming_hw: Union[Tuple[int, int], Dict[str, int]]=None):
        """
        Initialize the YOLO model with the given configuration.
        
        Args:
            config (Union[str, YOLOOptions]): Configuration object or YAML file path containing model parameters
            streaming_hw: Hardware dimensions for streaming mode configuration

        Setup model and complete configuration processing.        
        Handles:
        - Loading model to get class_names and model-specific info
        - Initializing class_names and nc from model
        - Handling pose task specific settings
        - Validating version
        - Storing task and pipeline references
        - Building preprocess and postprocess pipelines
        - Loading the model backend
        """
        self.config = init_yolo_config(config, streaming_hw)
        self.model = AutoBackend(self.config)

        self.config = init_model_info(self.config, self.model)

        self.task = self.config['task']
        self.class_names = self.config['class_names']
        self.classes = self.config['classes']
        self.conf = self.config['conf']
        self.iou = self.config['iou']
        self.streaming_mode = self.config['streaming_mode']
        self.img_shape = self.config['orig_imgs_shapes'] if self.streaming_mode else []

        self.preprocess = build_preprocess_pipeline(self.config)
        self.postprocess = build_postprocess_pipeline(self.config)

    def __call__(self, imgs: List[np.ndarray], dynamic_hw: Tuple = None) -> List[Union[Dict[str, Any], np.ndarray]]:
        """
        Perform inference on input images.
        
        Executes the complete inference pipeline: preprocess → model → postprocess
        
        Args:
            imgs: List of input images as numpy arrays
            
        Returns:
            List of prediction results, format depends on task type
        """
        pre_imgs = self.preprocess(imgs)
        preds = self.model(pre_imgs)
        shapes = get_imgs_shapes(
            imgs, 
            img_shape=self.img_shape, 
            streaming_mode=self.streaming_mode, 
            dynamic_hw=dynamic_hw
        )
        return self.postprocess(preds, shapes)

    def set_preprocess_params(self, **kwargs):
        """
        Dynamically update preprocessing parameters for speed.
        
        Usage Examples:
            # 1. Temporary override
            model.set_preprocess_params(pipeline="letterbox", size=640)
            
            # 2. Restore to original config settings
            model.set_preprocess_params(pipeline="original", size="original")
        """
        self.preprocess.update(**kwargs)

    def set_postprocess_params(self, **kwargs):
        """
        Dynamically update postprocessing parameters for speed.
        
        Usage Examples:
            # 1. Temporary override
            model.set_postprocess_params(conf=0.6, iou=0.5, pipeline="scale")
            
            # 2. Restore to original config settings
            model.set_postprocess_params(conf="original", iou="original", pipeline="original")
        """
        self.postprocess.update(**kwargs)


    @property
    def engine(self):
        """
        Access the underlying model backend.
        
        Returns:
            The model backend instance
        """
        return self.model



class YOLOToolBox(YOLO):
    """    
    This class extends the base YOLO class with additional utility methods
    for model evaluation, export, and performance testing.
    """
    def __init__(self, config: Union[str, YOLOOptions]):
        super().__init__(config)

    def __call__(self, imgs: List[np.ndarray], img_paths: List[str] = None) -> List[Any]:
        """
        Perform inference and construct results using the Results object.
        
        This method overrides the parent's __call__ method to provide
        additional result processing and formatting.
        
        Args:
            imgs: List of input images as numpy arrays
            img_paths: List of image file paths
            
        Returns:
            List of Results objects containing processed predictions
        """
        # Call parent's __call__ to get predictions
        preds = super().__call__(imgs)
        if img_paths is None:
            return preds
        return self.__construct_results(preds, imgs, img_paths)
    
    def export(
        self,
        model: Optional[Any] = None,
        shape: Tuple[int, int, int, int] = (1, 3, 224, 224),
        dynamic: bool = False,
        dynamic_batch: int = 16,
        precision: str = 'fp16',
        output_path: str = 'model.onnx',
        task: str = 'torch2engine'
    ) -> None:
        """
        Export the model to different formats (ONNX, TensorRT engine, etc.).
        
        Args:
            model: Model to export (defaults to the predictor's model)
            shape: Input shape for the model (batch_size, channels, height, width)
            onnx_path: If task is onnx2engine, please input onnx path 
            dynamic: Whether to use dynamic batch size
            precision: Precision type for export (fp16, fp32, etc.)
            output_path: Path where the exported model will be saved
            task: Export task type (torch2engine, torch2onnx, onnx2engine)
            model_wrapper: Optional wrapper function to modify the model before export
        """
        def ExportWrapper(model):
            """
            Factory function that creates an export wrapper for the model.
            Uses lazy import to avoid loading torch.nn until needed.
            """
            import torch.nn as nn
            
            class _ExportWrapper(nn.Module):
                def __init__(self, model):
                    super().__init__()
                    self.model = model.eval()

                def forward(self, x):
                    outputs = self.model(x)
                    return outputs[0]  # Only return detection tensor
            
            return _ExportWrapper(model)


        if self.config['weights'].split('.')[-1] == 'pt':
            import torch
            yolo_task = self.config['task']
            if model is None:
                model = self.model.model
                model.to(torch.device(self.config['device']))

            if yolo_task == 'segment':
                def custom_forward(self, x):
                    outputs = self.model(x)
                    return outputs[0], outputs[1][-1]
            else:
                def custom_forward(self, x):
                    outputs = self.model(x)
                    return outputs[0]
 
            ExportWrapper.forward = custom_forward

            model = ExportWrapper(model)       

        # if self.config['weights'].split('.')[-1] != 'pt':
        #     raise ValueError("Only pt weights are supported for export")
        if self.config['weights'].split('.')[-1] == 'onnx':
            onnx_path = self.config['weights']
            
        # Validate input shape format
        assert len(shape) == 4, "The input shape must be (1, C, H, W)"
        
        # Perform the specified export task
        if task == 'torch2engine':
            from nxva.nxengine.trt.convert import Convert        
            if onnx_path.split('.')[-1] != 'onnx':
                raise ValueError("Weights must be in ONNX format for onnx2engine task")
            Convert.torch_to_engine(model, shape, onnx_output_path=output_path, dynamic=dynamic, dynamic_batch=dynamic_batch, type=precision)
        elif task == 'torch2onnx':
            from nxva.nxengine.trt.convert import Convert
            if onnx_path.split('.')[-1] != 'pt':
                raise ValueError("Weights must be in PT format for torch2onnx task")
            Convert.torch_to_onnx(model, shape, onnx_output_path=output_path, dynamic=dynamic)
        elif task == 'onnx2engine':
            from nxva.nxengine.trt.convert import Convert
            if onnx_path.split('.')[-1] != 'onnx':
                raise ValueError("Weights must be in ONNX format for onnx2engine task")
            Convert.onnx_to_engine(onnx_path, shape, dynamic=dynamic, dynamic_batch=dynamic_batch, type=precision)
        # elif task == 'onnx2hef':
        # elif task == 'onnx2axcl':
            

        # elif task == 'onnx2jit':
        #     Convert.torch_to_jit(self.model, (1, 3, 224, 224), onnx_output_path=output_path, dynamic=False, type='fp16')
        else:
            raise ValueError(f"Unsupported task: {task}. Supported tasks: torch2engine, torch2onnx, onnx2engine")

    def speed_test(self, val_txt_path: str, batch_size: int = 1) -> None:
        """
        Perform speed testing on the model using validation data.
        
        This method measures the inference time for preprocessing, inference,
        and postprocessing stages separately.
        
        Args:
            val_txt_path: Path to validation data text file
            batch_size: Batch size for testing
            
        Returns:
            None (prints speed test results)
        """
        from nxva.toolbox.val import SpeedCalculator, Profile, DatasetLoader
        
        # Initialize speed calculator
        self.speed = SpeedCalculator()            
        
        # Create data loader for validation
        kpt_shape = self.config.get('kpt_shape', (17, 3))
        val_loader = DatasetLoader(val_txt_path, task=self.task, num_kpts=kpt_shape, batch_size=batch_size)

        # Helper to get orig_shapes
        def get_orig_shapes(imgs):
            if self.config['streaming_mode']:
                return self.config['orig_imgs_shapes']
            return [img.shape[:2] for img in imgs]

        #warmup
        for val_img, gt_label in val_loader:
            pre_imgs = self.preprocess(val_img)
            pred = self.model(pre_imgs)
            dets = self.postprocess(pred, get_orig_shapes(val_img))

        # Iterate through validation data and measure performance
        for val_img, gt_label in val_loader:
            
            # Create profilers for different stages
            profilers = (
                Profile(device=self.model.device),  # Preprocessing
                Profile(device=self.model.device),  # Inference
                Profile(device=self.model.device),  # Postprocessing
            )
            
            # Measure preprocessing time
            with profilers[0]:
                pre_imgs = self.preprocess(val_img)
                
            # Measure inference time
            with profilers[1]:
                pred = self.model(pre_imgs)
                
            # Measure postprocessing time
            with profilers[2]:
                dets = self.postprocess(pred, get_orig_shapes(val_img))
                
            # Update speed calculator with measurements
            self.speed.update(len(val_img), profilers)
            
        # Print speed test results
        print(self.speed.compute())

    def val(
        self,
        val_txt_path: str,
        batch_size: int = 1,
        save_dir: str = ".",
        plots: bool = False,
        conf: float = 0.25,
        iou: float = 0.5,
        verbose: bool = True,
    ) -> dict:
        """
        Run validation on the given val set and compute task metrics (mAP, confusion matrix, etc.).

        Args:
            val_txt_path: Path to a .txt file listing validation image paths (one per line).
            batch_size: Batch size for the validation loader.
            save_dir: Directory to save plots and confusion matrix images.
            plots: Whether to save PR curves and confusion matrix plots.
            conf: Confidence threshold for detection/segment/pose.
            iou: IoU threshold for NMS and matching.
            verbose: Whether to print per-class results.

        Returns:
            dict: Task metrics (e.g. precision, recall, mAP50, mAP50-95 for detect; top1_acc, top5_acc for classify).
        """
        from pathlib import Path
        size = self.config.get("size", 640)
        img_size = (size, size) if isinstance(size, int) else tuple(size)
        names = self.config.get("class_names") or {}
        nc = len(names) if names else self.config.get('nc', 80)
        common = dict(
            model=self,
            val_txt_path=val_txt_path,
            batch_size=batch_size,
            save_dir=Path(save_dir),
            plots=plots,
            conf=conf,
            iou=iou,
            names=names,
            nc=nc,
            img_size=img_size,
            verbose=verbose,
        )
        task = self.config.get("task", "detect")
        if task == "detect":
            from nxva.toolbox.val import DetectionValidator
            validator = DetectionValidator(**common)
        elif task == "segment":
            from nxva.toolbox.val import SegmentValidator
            validator = SegmentValidator(**common)
        elif task == "pose":
            from nxva.toolbox.val import PoseValidator
            kpt_shape = self.config.get("kpt_shape", (17, 3))
            validator = PoseValidator(**common, num_kpts=kpt_shape[0])
        elif task == "classify":
            from nxva.toolbox.val import ClassifyValidator
            validator = ClassifyValidator(
                model=self,
                val_txt_path=val_txt_path,
                batch_size=batch_size,
                save_dir=common["save_dir"],
                plots=plots,
                names=common["names"],
                nc=nc,
                img_size=img_size,
                verbose=verbose,
            )
        else:
            raise ValueError(f"Validation not implemented for task={task}. Use detect/segment/pose/classify.")
        return validator()

    def plot_results(self, results: List[Any], task: str) -> None:
        """
        Plot and save detection results.
        
        Args:
            results: List of Results objects
            task: Task name for output directory
            
        Returns:
            None
        """
        for i, result in enumerate(results):
            # create folder
            os.makedirs(f'./{task}_results', exist_ok=True)

            # save label as txt
            os.makedirs(f'./{task}_results/labels', exist_ok=True)
            result.save_txt(f'./{task}_results/labels/{task}_{i}.txt')

            # save image 
            os.makedirs(f'./{task}_results/images', exist_ok=True)
            img = result.plot(save=True, filename=f'./{task}_results/images/{task}_{i}.jpeg')

    def __construct_results(
        self,
        preds: List[Union[Dict[str, Any], np.ndarray]],
        orig_imgs: List[np.ndarray],
        img_paths: List[str]
    ) -> List[Any]:
        """
        Construct Results objects to encapsulate prediction results.
        
        This method uses the ultralytics Results object to standardize prediction
        results, enabling consistent visualization and processing across different
        YOLO tasks.
        
        Args:
            preds: List of prediction results from the model
            orig_imgs: List of original images in numpy array format
            img_paths: List of image file paths
            
        Returns:
            List of Results objects containing processed predictions
        """
        import torch
        from nxva.toolbox.result import Results
        
        results = []
        
        class_names = self.config['class_names']

        # Process predictions for each image based on task type
        for pred, orig_img, img_path in zip(preds, orig_imgs, img_paths):
            if self.task == 'detect':
                # Handle detection task - concatenate predictions if multiple detections
                if pred.shape[0] == 0:
                    pred = np.zeros((1, 6))
                pred = torch.from_numpy(pred)
                results.append(Results(orig_img=orig_img, path=img_path, names=class_names, boxes=pred[:, :6]))            
                
            elif self.task == 'pose':
                # Handle pose estimation task - separate boxes and keypoints
                boxes, kpts = pred['boxes'], pred['keypoints']
                r = Results(orig_img=orig_img, path=img_path, names=class_names, boxes=boxes)            
                r.update(keypoints=kpts)
                results.append(r)

            elif self.task == 'segment':
                # Handle segmentation task - filter predictions with valid masks
                box, masks = pred['boxes'], pred['mask']
                keep = masks.sum((-2, -1)) > 0  # only keep predictions with masks
                box, masks = box[keep], masks[keep]
                results.append(Results(orig_img=orig_img, path=img_path, names=class_names, boxes=box, masks=masks))

            elif self.task == 'classify':
                # Handle classification task - use probabilities
                results.append(Results(orig_img=orig_img, path=img_path, names=class_names, probs=pred[0]))  
        return results