#!/usr/bin/env python3
"""
vector_calibration_production.py

Production script for calibrating uncalibrated PIV vectors to physical units (m/s).
Converts pixel-based vectors to physical velocities using camera calibration models.
"""

import logging
import sys
from pathlib import Path

import cv2
import numpy as np
from scipy.io import loadmat, savemat

sys.path.append(str(Path(__file__).parent.parent))
from pivtools_core.paths import get_data_paths
from pivtools_core.vector_loading import load_coords_from_directory, read_mat_contents

# ===================== CONFIGURATION VARIABLES =====================
# Set these variables for your calibration setup
BASE_DIR = "/Users/morgan/Library/CloudStorage/OneDrive-UniversityofSouthampton/Documents/#current_processing/query_JHTDB/Planar_Images_with_wall/test"
image_count = 1000
DT_SECONDS = 1  # Time step between frames in seconds
CAMERA_NUM = 1  # Camera number (1-based)
MODEL_INDEX = 0  # Index of calibration model to use (0-based)
DOT_SPACING_MM = 28.89  # Physical spacing between calibration dots in mm
VECTOR_PATTERN = "%05d.mat"  # Pattern for vector files (e.g. "B%05d.mat", "%05d.mat")
TYPE_NAME = "instantaneous"  # Type name for uncalibrated data directory (e.g. "Instantaneous", "piv")
# Example: RUNS_TO_PROCESS = [1, 2, 3]  # Process only runs 1, 2, and 3
# ===================================================================

# Add src to path to import modules


# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class VectorCalibrator:
    def __init__(
        self,
        base_dir,
        camera_num,
        model_index,
        dt,
        dot_spacing_mm=28.89,
        vector_pattern="%05d.mat",
        type_name="Instantaneous",
    ):
        """
        Initialize vector calibrator

        Args:
            base_dir: Base directory containing data
            camera_num: Camera number (1-based)
            model_index: Index of calibration model to use (0-based)
            dt: Time step between frames in seconds
            dot_spacing_mm: Physical spacing of calibration dots in mm
            vector_pattern: Pattern for vector files (e.g. "B%05d.mat", "%05d.mat")
            type_name: Type name for uncalibrated data directory (e.g. "Instantaneous", "piv")
        """
        self.base_dir = Path(base_dir)
        self.camera_num = camera_num
        self.model_index = model_index
        self.dt = dt
        self.dot_spacing_mm = dot_spacing_mm
        self.vector_pattern = vector_pattern
        self.type_name = type_name

        # Load calibration model
        self.calibration_model = self._load_calibration_model()
        self.homography = self.calibration_model["homography"]
        self.camera_matrix = self.calibration_model["camera_matrix"]
        self.dist_coeffs = self.calibration_model["dist_coeffs"]
        self.dot_spacing_mm = dot_spacing_mm

        logger.info(f"Initialized calibrator for Camera {camera_num}")
        logger.info(f"Using calibration model index {model_index}")
        logger.info(f"Time step: {dt} seconds")
        logger.info(f"Dot spacing: {dot_spacing_mm} mm")
        logger.info(f"Vector pattern: {vector_pattern}")
        logger.info(f"Type name: {type_name}")

    def _load_calibration_model(self):
        """Load the specified calibration model"""
        calib_paths = get_data_paths(
            self.base_dir,
            num_images=1,  # Not used for calibration paths
            cam=self.camera_num,
            type_name="",  # Not used for calibration paths
            calibration=True,
        )

        calib_dir = calib_paths["calib_dir"]
        # Try common model filenames produced by planar calibrator
        # Always look for the model at "model/camera_model.mat"
        model_path = calib_dir / "model" / "camera_model.mat"
        if not model_path.exists():
            raise FileNotFoundError(f"Calibration model not found: {model_path}")

        logger.info(f"Loading calibration model: {model_path}")
        model_data = loadmat(str(model_path), squeeze_me=True, struct_as_record=False)

        required_fields = ["homography", "camera_matrix", "dist_coeffs"]
        missing_fields = [field for field in required_fields if field not in model_data]
        if missing_fields:
            raise ValueError(
                f"Missing required fields in calibration model: {missing_fields}"
            )

        # Ensure homography is 3x3
        homography = np.array(model_data["homography"])
        if homography.shape != (3, 3):
            raise ValueError(f"Homography must be 3x3, got shape {homography.shape}")

        # Ensure dist_coeffs is 1D array
        dist_coeffs = np.array(model_data["dist_coeffs"]).flatten()

        model_data["homography"] = homography.astype(np.float32)
        model_data["dist_coeffs"] = dist_coeffs.astype(np.float32)

        # Use dt from model if available, otherwise use instance dt
        if "dt" in model_data:
            logger.info(f"Using dt from calibration model: {model_data['dt']} seconds")
            self.dt = float(model_data["dt"])
        else:
            logger.info(
                f"No dt in calibration model, using provided dt: {self.dt} seconds"
            )

        return model_data

    def calibrate_coordinates(self, coords_x, coords_y):
        """
        Convert pixel coordinates to physical coordinates in mm using homography and distortion correction.

        Args:
            coords_x, coords_y: Coordinate arrays in pixels

        Returns:
            (x_mm, y_mm): Coordinate arrays in mm
        """
        # Stack coordinates for transformation
        pts = np.stack([coords_x.flatten(), coords_y.flatten()], axis=-1).astype(
            np.float32
        )

        # Ensure we have valid points
        if pts.size == 0:
            return coords_x, coords_y

        # Undistort points if distortion coefficients are present
        if np.any(self.dist_coeffs):
            # Reshape for cv2.undistortPoints: (N, 1, 2)
            pts_reshaped = pts.reshape(-1, 1, 2)
            try:
                pts_ud = cv2.undistortPoints(
                    pts_reshaped,
                    self.camera_matrix,
                    self.dist_coeffs,
                    P=self.camera_matrix,
                )
                pts_ud = pts_ud.reshape(-1, 2)
            except cv2.error as e:
                logger.warning(f"Undistortion failed, using original points: {e}")
                pts_ud = pts
        else:
            pts_ud = pts

        # Apply homography using OpenCV for numerical stability
        pts_ud = pts_ud.reshape(1, -1, 2)  # shape (1, N, 2) for perspectiveTransform
        pts_mapped = cv2.perspectiveTransform(pts_ud, self.homography)[0]

        # Reshape back to original shape
        x_mm = pts_mapped[:, 0].reshape(coords_x.shape)
        y_mm = pts_mapped[:, 1].reshape(coords_y.shape)

        logger.info(
            "Converted coordinates from pixels to mm (undistorted + homography)"
        )

        return x_mm, y_mm

    def calibrate_vectors(self, ux_px, uy_px, coords_x_px, coords_y_px):
        """
        Convert pixel-based velocity vectors to m/s using camera matrix (pinhole model) and distortion correction.

        Args:
            ux_px, uy_px: Velocity components in pixels/frame
            coords_x_px, coords_y_px: Grid coordinates in pixels

        Returns:
            (ux_ms, uy_ms): Velocity components in m/s
        """
        # Stack coordinates
        coords_px = np.stack([coords_x_px, coords_y_px], axis=-1).astype(np.float32)
        shape = coords_px.shape[:-1]
        coords_flat = coords_px.reshape(-1, 2)

        # Check for valid data
        if coords_flat.size == 0 or ux_px.size == 0 or uy_px.size == 0:
            logger.warning("Empty coordinate or vector data, returning zeros")
            return np.zeros_like(ux_px), np.zeros_like(uy_px)

        # Ensure arrays have compatible shapes
        if ux_px.shape != uy_px.shape or ux_px.shape != coords_x_px.shape:
            logger.error(
                f"Shape mismatch: ux_px={ux_px.shape}, uy_px={uy_px.shape}, coords_x_px={coords_x_px.shape}"
            )
            return np.zeros_like(ux_px), np.zeros_like(uy_px)

        # Undistort grid points
        if np.any(self.dist_coeffs):
            # Reshape for cv2.undistortPoints: (N, 1, 2)
            coords_reshaped = coords_flat.reshape(-1, 1, 2)
            try:
                coords_ud = cv2.undistortPoints(
                    coords_reshaped,
                    self.camera_matrix,
                    self.dist_coeffs,
                    P=self.camera_matrix,
                )
                coords_ud = coords_ud.reshape(-1, 2)
            except cv2.error as e:
                logger.warning(f"Grid undistortion failed, using original points: {e}")
                coords_ud = coords_flat
        else:
            coords_ud = coords_flat

        # Apply homography
        coords_ud = coords_ud.reshape(1, -1, 2)
        coords_mm = cv2.perspectiveTransform(coords_ud, self.homography)[0]

        # Displaced points
        disp_px = coords_flat + np.stack([ux_px.flatten(), uy_px.flatten()], axis=-1)

        if np.any(self.dist_coeffs):
            # Reshape for cv2.undistortPoints: (N, 1, 2)
            disp_reshaped = disp_px.reshape(-1, 1, 2)
            try:
                disp_ud = cv2.undistortPoints(
                    disp_reshaped,
                    self.camera_matrix,
                    self.dist_coeffs,
                    P=self.camera_matrix,
                )
                disp_ud = disp_ud.reshape(-1, 2)
            except cv2.error as e:
                logger.warning(
                    f"Displacement undistortion failed, using original points: {e}"
                )
                disp_ud = disp_px
        else:
            disp_ud = disp_px

        disp_ud = disp_ud.reshape(1, -1, 2)
        disp_mm = cv2.perspectiveTransform(disp_ud, self.homography)[0]

        # Metric displacement
        delta_mm = disp_mm - coords_mm

        # Convert to m/s
        ux_ms = (delta_mm[:, 0] / 1000.0) / self.dt  # mm to m, frame to s
        uy_ms = (delta_mm[:, 1] / 1000.0) / self.dt

        ux_ms = ux_ms.reshape(shape)
        uy_ms = uy_ms.reshape(shape)

        logger.info(
            "Converted vectors from pixels/frame to m/s using undistortion + homography"
        )

        return ux_ms, uy_ms

    def process_run(self, image_count, progress_cb=None):
        """
        Process and calibrate vectors for all available runs

        Args:
            image_count: Number of images in the run
            progress_cb: Optional callback for progress updates
        """
        logger.info(f"Processing run with {image_count} images")

        # Get data paths for uncalibrated data - use configured type
        paths = get_data_paths(
            self.base_dir,
            num_images=image_count,
            cam=self.camera_num,
            type_name=self.type_name,
            use_uncalibrated=True,
        )

        uncalib_data_dir = paths["data_dir"]
        logger.info(f"Uncalibrated data directory: {uncalib_data_dir}")

        # Get output paths for calibrated data
        calib_paths = get_data_paths(
            self.base_dir,
            num_images=image_count,
            cam=self.camera_num,
            type_name=self.type_name,
        )

        calib_data_dir = calib_paths["data_dir"]
        calib_data_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Calibrated data directory: {calib_data_dir}")

        if not uncalib_data_dir.exists():
            raise FileNotFoundError(
                f"Uncalibrated data directory not found: {uncalib_data_dir}"
            )

        # Load coordinates for all runs
        logger.info("Loading coordinates...")
        x_coords_list, y_coords_list = load_coords_from_directory(
            uncalib_data_dir, runs=None  # Load all available runs
        )

        if not x_coords_list:
            logger.error("No coordinate data found!")
            raise ValueError("No coordinate data found")

        logger.info(f"Loaded coordinates for {len(x_coords_list)} runs")

        # Find runs with valid data
        valid_runs = []
        for i, (x_coords, y_coords) in enumerate(zip(x_coords_list, y_coords_list)):
            run_num = i + 1
            # Ensure None is replaced with empty arrays
            if x_coords is None:
                x_coords = np.array([])
            if y_coords is None:
                y_coords = np.array([])
            valid_coords = np.sum(~np.isnan(x_coords)) + np.sum(~np.isnan(y_coords))
            logger.info(f"Run {run_num}: {valid_coords} valid coordinates")
            if valid_coords > 0:
                valid_runs.append((i, run_num, valid_coords))

        if not valid_runs:
            raise ValueError("No runs with valid coordinate data found")

        logger.info(
            f"Found {len(valid_runs)} runs with valid data: {[r[1] for r in valid_runs]}"
        )

        # Create coordinate structure
        max_run = max([r[1] for r in valid_runs])
        coord_dtype = np.dtype([("x", "O"), ("y", "O")])
        coordinates = np.empty(max_run, dtype=coord_dtype)

        # Process each valid run
        for run_idx, run_num, valid_coord_count in valid_runs:
            logger.info(
                f"Processing run {run_num} with {valid_coord_count} valid coordinates"
            )
            x_coords_px = x_coords_list[run_idx]
            y_coords_px = y_coords_list[run_idx]

            # Calibrate coordinates to mm
            x_coords_mm, y_coords_mm = self.calibrate_coordinates(
                x_coords_px, y_coords_px
            )
            coordinates[run_num - 1] = (x_coords_mm, y_coords_mm)

        # Fill all runs: valid runs get data, others get empty arrays
        valid_run_indices = set(r[1] for r in valid_runs)
        for run_num in range(1, max_run + 1):
            if run_num in valid_run_indices:
                # Already set above for valid runs
                continue
            coordinates[run_num - 1] = (np.array([]), np.array([]))

        coords_output = {"coordinates": coordinates}

        # Save calibrated coordinates
        calibrated_dir = (
            self.base_dir
            / "calibrated_piv"
            / str(image_count)
            / f"Cam{self.camera_num}"
            / self.type_name
        )
        calibrated_dir.mkdir(parents=True, exist_ok=True)
        coords_path = calibrated_dir / "coordinates.mat"
        savemat(str(coords_path), coords_output)
        logger.info(f"Saved calibrated coordinates: {coords_path}")

        # Process vector files using the first valid run's coordinates
        if valid_runs:
            first_run_idx = valid_runs[0][0]
            x_coords_for_vectors = x_coords_list[first_run_idx]
            y_coords_for_vectors = y_coords_list[first_run_idx]

            self._process_vector_files(
                uncalib_data_dir,
                calib_data_dir,
                image_count,
                x_coords_for_vectors,
                y_coords_for_vectors,
                max_run,
                valid_runs,
                progress_cb,
            )
        else:
            logger.error("No valid runs found for vector processing")

    def _process_vector_files(
        self,
        uncalib_dir,
        calib_dir,
        num_images,
        coords_x_px,
        coords_y_px,
        max_run,
        valid_runs,
        progress_cb,
    ):
        """Process all vector files in the directory"""
        logger.info("Processing vector files...")

        # Use configured vector pattern
        vector_pattern = self.vector_pattern

        processed_vectors = []

        for i in range(1, num_images + 1):
            vector_file = uncalib_dir / (vector_pattern % i)

            if not vector_file.exists():
                if i <= 5:  # Only log first few missing files
                    logger.warning(f"Vector file not found: {vector_file}")
                continue

            try:
                # Load uncalibrated vectors
                vector_data = read_mat_contents(str(vector_file))  # Shape varies

                # Handle different vector data formats
                if vector_data.ndim == 4 and vector_data.shape[0] == 1:
                    # Single run format: (1, 3, H, W)
                    ux_px = vector_data[0, 0, :, :]
                    uy_px = vector_data[0, 1, :, :]
                    b_mask = vector_data[0, 2, :, :]
                elif vector_data.ndim == 3 and vector_data.shape[0] == 3:
                    # Single run format: (3, H, W)
                    ux_px = vector_data[0, :, :]
                    uy_px = vector_data[1, :, :]
                    b_mask = vector_data[2, :, :]
                else:
                    logger.warning(
                        f"Unexpected vector data shape in {vector_file.name}: {vector_data.shape}"
                    )
                    continue

                # Calibrate vectors
                ux_ms, uy_ms = self.calibrate_vectors(
                    ux_px, uy_px, coords_x_px, coords_y_px
                )

                # Create piv_result structure array with proper MATLAB struct format
                piv_dtype = np.dtype([("ux", "O"), ("uy", "O"), ("b_mask", "O")])
                piv_result = np.empty(max_run, dtype=piv_dtype)

                for run_num in range(1, max_run + 1):
                    run_idx = run_num - 1  # Convert to 0-based

                    # Check if this run has valid data
                    run_has_data = any(r[1] == run_num for r in valid_runs)

                    if run_has_data:
                        # This creates piv_result(run_idx).ux, piv_result(run_idx).uy, piv_result(run_idx).b_mask
                        piv_result[run_idx] = (ux_ms, uy_ms, b_mask)
                    else:
                        # Empty struct for runs not being processed
                        piv_result[run_idx] = (np.array([]), np.array([]), np.array([]))

                # Save calibrated piv_result into calibrated_piv output tree
                output_file = calib_dir / (self.vector_pattern % i)
                savemat(str(output_file), {"piv_result": piv_result})

                processed_vectors.append(
                    {"ux_ms": ux_ms, "uy_ms": uy_ms, "b_mask": b_mask, "frame": i}
                )

                # Progress callback
                if progress_cb:
                    progress = (i / num_images) * 100
                    progress_cb(
                        {
                            "processed_frames": i,
                            "total_frames": num_images,
                            "progress": progress,
                            "successful_frames": len(processed_vectors),
                        }
                    )

            except Exception as e:
                logger.error(f"Failed to process {vector_file.name}: {str(e)}")
                continue

        logger.info(
            f"Successfully processed {len(processed_vectors)} vector files into {calib_dir}"
        )


def main():
    logger.info("Starting vector calibration with configuration:")
    logger.info(f"Base directory: {BASE_DIR}")
    logger.info(f"Run number: {image_count}")
    logger.info(f"Time step: {DT_SECONDS} seconds")
    logger.info(f"Camera: {CAMERA_NUM}")
    logger.info(f"Model index: {MODEL_INDEX}")
    logger.info(f"Dot spacing: {DOT_SPACING_MM} mm")
    logger.info(f"Vector pattern: {VECTOR_PATTERN}")
    logger.info(f"Type name: {TYPE_NAME}")

    try:
        calibrator = VectorCalibrator(
            base_dir=BASE_DIR,
            camera_num=CAMERA_NUM,
            model_index=MODEL_INDEX,
            dt=DT_SECONDS,
            dot_spacing_mm=DOT_SPACING_MM,
            vector_pattern=VECTOR_PATTERN,
            type_name=TYPE_NAME,
        )

        calibrator.process_run(image_count)

        logger.info("Vector calibration completed successfully")

    except Exception as e:
        logger.error(f"Calibration failed: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()
