"""
Touchscreen calibration module.

This module provides calibration functionality to align touchscreen coordinates
with display pixels. Uses a multi-point calibration approach with affine transformation.

The calibration process:
1. Display calibration targets (circles) at known display positions
2. User touches each target
3. Record both display coordinates and touch coordinates
4. Calculate affine transformation matrix to map touch -> display
5. Save calibration data for future use
"""

import json
import math
from dataclasses import dataclass, asdict
from typing import List, Tuple, Optional
from pathlib import Path


@dataclass
class CalibrationPoint:
    """
    A single calibration point pair.

    Attributes:
        display_x: X coordinate on display (pixels)
        display_y: Y coordinate on display (pixels)
        touch_x: Raw X coordinate from touch sensor
        touch_y: Raw Y coordinate from touch sensor
    """
    display_x: int
    display_y: int
    touch_x: int
    touch_y: int


@dataclass
class CalibrationData:
    """
    Complete calibration dataset with transformation matrix.

    Attributes:
        points: List of calibration point pairs
        matrix: 2x3 affine transformation matrix [a, b, c, d, e, f]
                where: x' = ax + by + c, y' = dx + ey + f
        width: Display width in pixels
        height: Display height in pixels
        rms_error: Root mean square error of calibration (pixels)
    """
    points: List[CalibrationPoint]
    matrix: List[float]  # [a, b, c, d, e, f]
    width: int
    height: int
    rms_error: float = 0.0

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            'points': [asdict(p) for p in self.points],
            'matrix': self.matrix,
            'width': self.width,
            'height': self.height,
            'rms_error': self.rms_error,
        }

    @classmethod
    def from_dict(cls, data: dict) -> 'CalibrationData':
        """Create from dictionary (JSON deserialization)."""
        points = [CalibrationPoint(**p) for p in data['points']]
        return cls(
            points=points,
            matrix=data['matrix'],
            width=data['width'],
            height=data['height'],
            rms_error=data.get('rms_error', 0.0),
        )


class TouchCalibration:
    """
    Touchscreen calibration with affine transformation.

    This class handles:
    - Generating calibration target positions
    - Computing affine transformation from calibration points
    - Transforming raw touch coordinates to display coordinates
    - Saving/loading calibration data

    Args:
        width: Display width in pixels
        height: Display height in pixels
        num_points: Number of calibration points (default 9 for 3x3 grid)
    """

    def __init__(self, width: int, height: int, num_points: int = 9):
        self.width = width
        self.height = height
        self.num_points = num_points
        self.calibration_data: Optional[CalibrationData] = None

    def generate_target_positions(self, margin: int = 100, target_radius: int = 10) -> List[Tuple[int, int]]:
        """
        Generate positions for calibration targets.

        Creates a grid of targets with margins from edges.
        For 9 points: 3x3 grid (corners, edges, center)
        For 5 points: corners + center

        Args:
            margin: Distance from screen edges (pixels)
            target_radius: Radius of calibration circles (pixels)

        Returns:
            List of (x, y) positions for calibration targets
        """
        if self.num_points == 5:
            # 5-point calibration: corners + center
            return [
                (margin, margin),                           # Top-left
                (self.width - margin, margin),              # Top-right
                (self.width - margin, self.height - margin), # Bottom-right
                (margin, self.height - margin),             # Bottom-left
                (self.width // 2, self.height // 2),        # Center
            ]
        elif self.num_points == 9:
            # 9-point calibration: 3x3 grid
            mid_x = self.width // 2
            mid_y = self.height // 2
            return [
                (margin, margin),                           # Top-left
                (mid_x, margin),                            # Top-center
                (self.width - margin, margin),              # Top-right
                (margin, mid_y),                            # Middle-left
                (mid_x, mid_y),                             # Center
                (self.width - margin, mid_y),               # Middle-right
                (margin, self.height - margin),             # Bottom-left
                (mid_x, self.height - margin),              # Bottom-center
                (self.width - margin, self.height - margin), # Bottom-right
            ]
        else:
            # Custom grid based on num_points
            # Create as uniform a grid as possible
            grid_size = int(math.sqrt(self.num_points))
            positions = []
            for i in range(grid_size):
                for j in range(grid_size):
                    x = margin + (self.width - 2 * margin) * j // (grid_size - 1)
                    y = margin + (self.height - 2 * margin) * i // (grid_size - 1)
                    positions.append((x, y))
            return positions[:self.num_points]

    def add_calibration_point(self, display_x: int, display_y: int,
                             touch_x: int, touch_y: int) -> None:
        """
        Add a calibration point pair.

        Args:
            display_x: X coordinate on display
            display_y: Y coordinate on display
            touch_x: Raw X from touch sensor
            touch_y: Raw Y from touch sensor
        """
        if self.calibration_data is None:
            self.calibration_data = CalibrationData(
                points=[],
                matrix=[1, 0, 0, 0, 1, 0],  # Identity transform
                width=self.width,
                height=self.height,
            )

        point = CalibrationPoint(display_x, display_y, touch_x, touch_y)
        self.calibration_data.points.append(point)

    def compute_calibration(self) -> bool:
        """
        Compute affine transformation matrix from calibration points.

        Uses least-squares fitting to find the best affine transformation
        that maps touch coordinates to display coordinates.

        Affine transformation:
            x_display = a * x_touch + b * y_touch + c
            y_display = d * x_touch + e * y_touch + f

        Returns:
            True if calibration successful, False otherwise
        """
        if not self.calibration_data or len(self.calibration_data.points) < 3:
            return False

        # Build least-squares system
        # For each point: x' = ax + by + c, y' = dx + ey + f
        # We need at least 3 points to solve for 6 unknowns

        points = self.calibration_data.points
        n = len(points)

        # Build matrices for least-squares: Ax = b
        # For x': [x1 y1 1] [a]   [x'1]
        #         [x2 y2 1] [b] = [x'2]
        #         [x3 y3 1] [c]   [x'3]

        # Matrix A (n x 3)
        A = [[p.touch_x, p.touch_y, 1] for p in points]

        # Vectors b_x and b_y
        b_x = [p.display_x for p in points]
        b_y = [p.display_y for p in points]

        # Solve least-squares for x-transform: [a, b, c]
        abc = self._solve_least_squares(A, b_x)
        if abc is None:
            return False

        # Solve least-squares for y-transform: [d, e, f]
        def_vals = self._solve_least_squares(A, b_y)
        if def_vals is None:
            return False

        # Store transformation matrix
        self.calibration_data.matrix = abc + def_vals

        # Compute RMS error
        self.calibration_data.rms_error = self._compute_rms_error()

        return True

    def _solve_least_squares(self, A: List[List[float]], b: List[float]) -> Optional[List[float]]:
        """
        Solve least-squares problem: A x = b

        Uses normal equations: (A^T A) x = A^T b

        Args:
            A: Matrix (n x 3)
            b: Vector (n x 1)

        Returns:
            Solution vector x (3 x 1) or None if singular
        """
        # Compute A^T A (3 x 3)
        n = len(A)
        m = len(A[0])

        ATA = [[0.0] * m for _ in range(m)]
        for i in range(m):
            for j in range(m):
                for k in range(n):
                    ATA[i][j] += A[k][i] * A[k][j]

        # Compute A^T b (3 x 1)
        ATb = [0.0] * m
        for i in range(m):
            for k in range(n):
                ATb[i] += A[k][i] * b[k]

        # Solve 3x3 system using Gaussian elimination
        return self._solve_3x3(ATA, ATb)

    def _solve_3x3(self, A: List[List[float]], b: List[float]) -> Optional[List[float]]:
        """
        Solve 3x3 linear system using Gaussian elimination.

        Args:
            A: 3x3 matrix
            b: 3x1 vector

        Returns:
            Solution vector or None if singular
        """
        # Create augmented matrix
        aug = [A[i][:] + [b[i]] for i in range(3)]

        # Forward elimination
        for i in range(3):
            # Find pivot
            max_row = i
            for k in range(i + 1, 3):
                if abs(aug[k][i]) > abs(aug[max_row][i]):
                    max_row = k

            # Swap rows
            aug[i], aug[max_row] = aug[max_row], aug[i]

            # Check for singular matrix
            if abs(aug[i][i]) < 1e-10:
                return None

            # Eliminate column
            for k in range(i + 1, 3):
                factor = aug[k][i] / aug[i][i]
                for j in range(i, 4):
                    aug[k][j] -= factor * aug[i][j]

        # Back substitution
        x = [0.0] * 3
        for i in range(2, -1, -1):
            x[i] = aug[i][3]
            for j in range(i + 1, 3):
                x[i] -= aug[i][j] * x[j]
            x[i] /= aug[i][i]

        return x

    def _compute_rms_error(self) -> float:
        """
        Compute root mean square error of calibration.

        Returns:
            RMS error in pixels
        """
        if not self.calibration_data or not self.calibration_data.points:
            return 0.0

        total_sq_error = 0.0
        for point in self.calibration_data.points:
            # Transform touch coordinates
            tx, ty = self.transform(point.touch_x, point.touch_y)

            # Compute error
            dx = tx - point.display_x
            dy = ty - point.display_y
            total_sq_error += dx * dx + dy * dy

        return math.sqrt(total_sq_error / len(self.calibration_data.points))

    def transform(self, touch_x: int, touch_y: int) -> Tuple[int, int]:
        """
        Transform raw touch coordinates to display coordinates.

        Args:
            touch_x: Raw X from touch sensor
            touch_y: Raw Y from touch sensor

        Returns:
            (display_x, display_y) tuple
        """
        if not self.calibration_data:
            # No calibration - return raw coordinates
            return (touch_x, touch_y)

        m = self.calibration_data.matrix
        a, b, c, d, e, f = m

        # Apply affine transformation
        x = a * touch_x + b * touch_y + c
        y = d * touch_x + e * touch_y + f

        # Clamp to display bounds
        x = max(0, min(int(round(x)), self.width - 1))
        y = max(0, min(int(round(y)), self.height - 1))

        return (x, y)

    def save(self, filepath: str) -> None:
        """
        Save calibration data to JSON file.

        Args:
            filepath: Path to save calibration file
        """
        if not self.calibration_data:
            raise ValueError("No calibration data to save")

        path = Path(filepath)
        path.parent.mkdir(parents=True, exist_ok=True)

        with open(filepath, 'w') as f:
            json.dump(self.calibration_data.to_dict(), f, indent=2)

    def load(self, filepath: str) -> bool:
        """
        Load calibration data from JSON file.

        Args:
            filepath: Path to calibration file

        Returns:
            True if loaded successfully, False otherwise
        """
        try:
            with open(filepath, 'r') as f:
                data = json.load(f)

            self.calibration_data = CalibrationData.from_dict(data)

            # Verify dimensions match
            if (self.calibration_data.width != self.width or
                self.calibration_data.height != self.height):
                print(f"Warning: Calibration dimensions ({self.calibration_data.width}x{self.calibration_data.height}) "
                      f"don't match display ({self.width}x{self.height})")
                return False

            return True
        except (FileNotFoundError, json.JSONDecodeError, KeyError) as e:
            print(f"Failed to load calibration: {e}")
            return False

    def is_calibrated(self) -> bool:
        """Check if calibration is loaded and valid."""
        return (self.calibration_data is not None and
                len(self.calibration_data.points) >= 3 and
                self.calibration_data.matrix is not None)

    def get_calibration_quality(self) -> str:
        """
        Get calibration quality assessment.

        Returns:
            Quality string: "Excellent", "Good", "Fair", "Poor", or "Uncalibrated"
        """
        if not self.is_calibrated():
            return "Uncalibrated"

        error = self.calibration_data.rms_error

        if error < 5:
            return "Excellent"
        elif error < 10:
            return "Good"
        elif error < 20:
            return "Fair"
        else:
            return "Poor"