"""Phase 1 transform pipeline: letterbox resize, normalize, horizontal flip.

Camera-generalization augmentation (homography warps, photometric jitter) is
explicitly deferred to Phase 2 per the project plan — this pipeline is the minimal
set needed to prove the core model trains.

All lane coordinates are carried as plain (x, y) pixel-space lists until the final
`ToSampledTargets` step, which normalizes to [0,1] and resamples onto a fixed y-grid
so batches of variable lane-count/length can be collated into fixed-size tensors.
"""
from __future__ import annotations

import random

import cv2
import numpy as np
import torch

from utils.curve import compute_letterbox, sample_lane_at_ys

IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, sample: dict) -> dict:
        for t in self.transforms:
            sample = t(sample)
        return sample


class LetterboxResize:
    def __init__(self, out_w: int, out_h: int):
        self.out_w = out_w
        self.out_h = out_h

    def __call__(self, sample: dict) -> dict:
        image = sample["image"]
        src_h, src_w = image.shape[:2]
        lb = compute_letterbox(src_w, src_h, self.out_w, self.out_h)

        new_w, new_h = int(round(src_w * lb.scale)), int(round(src_h * lb.scale))
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

        canvas = np.zeros((self.out_h, self.out_w, 3), dtype=image.dtype)
        px, py = int(round(lb.pad_x)), int(round(lb.pad_y))
        canvas[py:py + new_h, px:px + new_w] = resized

        lanes = []
        for lane in sample["lanes"]:
            pts = np.array(lane, dtype=np.float32)
            pts = lb.apply_points(pts)
            lanes.append([(float(x), float(y)) for x, y in pts])

        sample["image"] = canvas
        sample["lanes"] = lanes
        sample["letterbox"] = lb
        return sample


class RandomHorizontalFlip:
    def __init__(self, p: float = 0.5):
        self.p = p

    def __call__(self, sample: dict) -> dict:
        if random.random() >= self.p:
            return sample
        image = sample["image"]
        w = image.shape[1]
        sample["image"] = np.ascontiguousarray(image[:, ::-1, :])
        sample["lanes"] = [[(w - 1 - x, y) for (x, y) in lane] for lane in sample["lanes"]]
        return sample


class Normalize:
    """Uint8 HWC image -> normalized float32 CHW torch tensor."""

    def __call__(self, sample: dict) -> dict:
        image = sample["image"].astype(np.float32) / 255.0
        image = (image - IMAGENET_MEAN) / IMAGENET_STD
        sample["image"] = torch.from_numpy(image.transpose(2, 0, 1)).float()
        return sample


class ToSampledTargets:
    """Resample variable-length lane polylines onto a fixed-size training target.

    Produces, for a fixed grid of `num_sample_ys` normalized y-values shared by every
    sample in a batch:
      - xs:        (max_lanes, num_sample_ys) normalized x at each sample_y
      - valid_mask:(max_lanes, num_sample_ys) bool, True where the lane is defined
      - lane_valid:(max_lanes,) bool, True for real (non-padding) lane slots
      - endpoints: (max_lanes, 2) normalized (y_start, y_end)
    """

    def __init__(self, out_w: int, out_h: int, max_lanes: int, num_sample_ys: int):
        self.out_w = out_w
        self.out_h = out_h
        self.max_lanes = max_lanes
        self.sample_ys = np.linspace(0.0, 1.0, num_sample_ys, dtype=np.float32)

    def __call__(self, sample: dict) -> dict:
        lanes = sample["lanes"][: self.max_lanes]
        n = len(self.sample_ys)

        xs = np.zeros((self.max_lanes, n), dtype=np.float32)
        valid_mask = np.zeros((self.max_lanes, n), dtype=bool)
        lane_valid = np.zeros((self.max_lanes,), dtype=bool)
        endpoints = np.zeros((self.max_lanes, 2), dtype=np.float32)

        for i, lane in enumerate(lanes):
            if len(lane) < 2:
                continue
            norm_lane = [(x / self.out_w, y / self.out_h) for (x, y) in lane]
            lane_xs, lane_valid_mask = sample_lane_at_ys(norm_lane, self.sample_ys)
            if not lane_valid_mask.any():
                continue
            xs[i] = lane_xs
            valid_mask[i] = lane_valid_mask
            lane_valid[i] = True
            ys_in_lane = self.sample_ys[lane_valid_mask]
            endpoints[i] = [ys_in_lane.min(), ys_in_lane.max()]

        sample["target_xs"] = torch.from_numpy(xs)
        sample["target_valid_mask"] = torch.from_numpy(valid_mask)
        sample["target_lane_valid"] = torch.from_numpy(lane_valid)
        sample["target_endpoints"] = torch.from_numpy(endpoints)
        sample["sample_ys"] = torch.from_numpy(self.sample_ys)
        return sample


def build_transforms(out_w: int, out_h: int, max_lanes: int, num_sample_ys: int, train: bool) -> Compose:
    steps = [LetterboxResize(out_w, out_h)]
    if train:
        steps.append(RandomHorizontalFlip(p=0.5))
    steps.append(ToSampledTargets(out_w, out_h, max_lanes, num_sample_ys))
    steps.append(Normalize())
    return Compose(steps)