"""
LASER ASD Wrapper

Simple interface for LoCoNet + LASER Active Speaker Detection model.

Based on: https://github.com/plnguyen2908/LASER_ASD
"""

import logging
import os
import sys
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
import torch
import torch.nn as nn

# Add laser module to path
LASER_DIR = Path(__file__).parent
sys.path.insert(0, str(LASER_DIR))

# Setup logging
logger = logging.getLogger(__name__)

# Default model path (can be set via environment variable)
DEFAULT_MODEL_PATH = Path(
    os.environ.get(
        "LASER_ASD_MODEL_PATH",
        str(Path.home() / ".cache" / "laser_asd" / "loconet_laser.model")
    )
)


class LaserASDModel:
    """
    LASER ASD (Lip Landmark Assisted Speaker Detection) Model

    Provides a simple interface for the LoCoNet encoder with LASER
    landmark injection for Active Speaker Detection.

    Reference: https://github.com/plnguyen2908/LASER_ASD
    """

    def __init__(
        self,
        device: str = "cuda",
        model_path: Optional[Path] = None,
        use_landmarks: bool = False,
    ):
        """
        Initialize LASER ASD model.

        Args:
            device: Device to run model on ('cuda' or 'cpu')
            model_path: Path to model weights file
            use_landmarks: Whether to use landmark features (False uses zeros)
        """
        self.device = device
        self.model_path = model_path or DEFAULT_MODEL_PATH
        self.use_landmarks = use_landmarks

        self._model = None
        self._loaded = False

    def load_weights(self, model_path: Optional[str] = None):
        """
        Load model weights.

        Args:
            model_path: Optional path to model weights file
        """
        if model_path:
            self.model_path = Path(model_path)
        self._load_model()

    def _load_model(self):
        """Lazy load the model."""
        if self._loaded:
            return

        try:
            # Import LASER model components
            from .landmark_loconet import Loconet

            # Create config-like object matching LASER's multi.yaml
            class ModelConfig:
                NUM_SPEAKERS = 3
                CLIP_LENGTH = 200
                AV = "speaker_temporal"
                AV_layers = 3
                ADJUST_ATTENTION = 0
                AUDIO_MODEL = 'resnet18'
                VISUAL_MODEL = 'resnet18'

            class DataConfig:
                numWorkers = 4
                dataPathAVA = '/tmp'

            class Config:
                DATA = DataConfig()
                MODEL = ModelConfig()

            cfg = Config()

            # Initialize model (n_channel=4 for landmarks, layer=1 for injection point)
            self._model = Loconet(cfg, n_channel=4, layer=1)

            # Load weights
            if self.model_path.exists():
                state_dict = torch.load(str(self.model_path), map_location=self.device)
                # Handle different checkpoint formats
                if 'model_state_dict' in state_dict:
                    self._model.load_state_dict(state_dict['model_state_dict'], strict=False)
                else:
                    self._model.load_state_dict(state_dict, strict=False)
                logger.info(f"LASER model loaded from {self.model_path}")
            else:
                logger.warning(f"Model weights not found: {self.model_path}")

            self._model = self._model.to(self.device)
            self._model.eval()
            self._loaded = True

        except Exception as e:
            logger.error(f"Failed to load LASER model: {e}")
            raise

    def predict(
        self,
        face_crops: np.ndarray,
        audio_data: np.ndarray,
        sample_rate: int = 16000,
        fps: float = 25.0,
    ) -> np.ndarray:
        """
        Predict speaking probability for each frame.

        Args:
            face_crops: Face crop images [T, H, W, C] or [T, H, W]
            audio_data: Audio waveform [samples]
            sample_rate: Audio sample rate (default: 16000)
            fps: Video frame rate (default: 25.0)

        Returns:
            Per-frame speaking scores [T] (>= 0 means speaking)
        """
        self._load_model()

        # Prepare visual features
        visual_feature = self._prepare_visual_features(face_crops, fps)

        # Prepare audio features
        audio_feature = self._prepare_audio_features(audio_data, sample_rate, visual_feature.shape[2], fps)

        # Create empty landmarks
        b, s, t = visual_feature.shape[:3]
        landmarks = torch.zeros((b, s, t, 82, 2), device=self.device)

        with torch.no_grad():
            pred_score = self._model.model.forward_evaluation(
                audio_feature,
                visual_feature,
                landmarks,
                None,
                None,
                self.use_landmarks
            )

        return pred_score.cpu().numpy()

    def _prepare_visual_features(
        self,
        face_crops: np.ndarray,
        fps: float,
    ) -> torch.Tensor:
        """
        Convert face crops to visual features.

        Args:
            face_crops: [T, H, W, C] or [T, H, W]

        Returns:
            [B=1, S=3, T, 112, 112] tensor
        """
        T = len(face_crops)

        # Resize and convert to grayscale
        processed = []
        for frame in face_crops:
            if len(frame.shape) == 3:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            frame = cv2.resize(frame, (112, 112))
            processed.append(frame)

        # Stack and convert to tensor
        visual = np.stack(processed, axis=0)  # [T, 112, 112]
        visual = torch.from_numpy(visual).float().to(self.device)

        # Reshape to [B=1, S=3, T, 112, 112]
        # S=3 is for self + 2 context faces (we use self 3 times for simplicity)
        visual = visual.unsqueeze(0)  # [1, T, 112, 112]
        visual = visual.unsqueeze(0).repeat(1, 3, 1, 1, 1)  # [1, 3, T, 112, 112]

        return visual

    def _prepare_audio_features(
        self,
        audio_data: np.ndarray,
        sample_rate: int,
        num_frames: int,
        fps: float,
    ) -> torch.Tensor:
        """
        Convert audio data to MFCC features.

        Args:
            audio_data: [samples]
            sample_rate: Sampling rate
            num_frames: Number of frames
            fps: Frame rate

        Returns:
            [B=1, 1, 4*T, 64] tensor
        """
        from .torchvggish import vggish_input

        # Normalize audio
        if audio_data.dtype != np.float32:
            audio_data = audio_data.astype(np.float32)
        if np.abs(audio_data).max() > 1.0:
            audio_data = audio_data / 32768.0  # Assume 16-bit PCM

        # Extract MFCC features
        audio_feature = vggish_input.waveform_to_examples(
            audio_data, sample_rate, num_frames, fps, return_tensor=False
        )

        # Convert to tensor [B=1, 1, 4*T, 64]
        audio_tensor = torch.from_numpy(audio_feature).float().to(self.device)
        audio_tensor = audio_tensor.unsqueeze(0).unsqueeze(0)

        return audio_tensor


def create_laser_model(
    device: str = "cuda",
    model_path: Optional[str] = None,
    **kwargs,
) -> LaserASDModel:
    """
    Factory function to create LASER ASD model.

    Args:
        device: Device ('cuda' or 'cpu')
        model_path: Path to model weights

    Returns:
        LaserASDModel instance
    """
    path = Path(model_path) if model_path else None
    return LaserASDModel(device=device, model_path=path, **kwargs)
