"""RGB sensor with noise model and ISP."""

import json

import torch

from deeplens.sensor import Sensor
from deeplens.sensor.isp_modules.isp import InvertibleISP


class RGBSensor(Sensor):
    """RGB Bayer-pattern sensor with physics-based noise model and invertible ISP.

    Simulates the full image-capture pipeline from linear photon counts to
    display-ready sRGB:

    1. **Spectral integration** – optional per-channel spectral response.
    2. **Bayer mosaic** – pixel-level colour filtering to a single-channel
       raw image.
    3. **Noise** – shot noise (signal-dependent Gaussian) + read noise
       (ISO-independent Gaussian) added to the n-bit raw data.
    4. **ISP** (forward) – via an
       :class:`~deeplens.sensor.isp_modules.isp.InvertibleISP`:
       black-level correction → white balance → colour matrix →
       demosaicing → gamma correction.

    The ISP is **invertible**: ``unprocess()`` converts sRGB back to linear
    RGB for training data generation.

    Attributes:
        bit (int): ADC bit depth.
        nbit_max (int): Maximum digital number ``2**bit - 1``.
        black_level (int): Black level pedestal [DN].
        bayer_pattern (str): Bayer pattern (e.g. ``"rggb"``).
        iso_base (int): Base ISO (noise-free reference).
        isp (InvertibleISP): Embedded ISP pipeline.
    """

    def __init__(
        self,
        size=(36.0, 24.0),
        res=(5472, 3648),
        bit=10,
        black_level=64,
        bayer_pattern="rggb",
        white_balance=(2.0, 1.0, 1.8),
        color_matrix=None,
        gamma_param=2.2,
        iso_base=100,
        read_noise_std=0.5,
        shot_noise_std_alpha=0.4,
        shot_noise_std_beta=0.0,
        wavelengths=None,
        red_response=None,
        green_response=None,
        blue_response=None,
    ):
        """
        Args:
            size (tuple): Sensor physical size in mm (W, H). Default (36.0, 24.0).
            res (tuple): Sensor resolution in pixels (W, H). Default (5472, 3648).
            bit (int): Bit depth. Default 10.
            black_level (int): Black level. Default 64.
            bayer_pattern (str): Bayer pattern e.g. "rggb". Default "rggb".
            white_balance (tuple): White balance gains. Default (2.0, 1.0, 1.8).
            color_matrix (list or Tensor): Color correction matrix.
            gamma_param (float): Gamma correction parameter. Default 2.2.
            iso_base (int): Base ISO. Default 100.
            read_noise_std (float): Read noise std. Default 0.5.
            shot_noise_std_alpha (float): Shot noise alpha. Default 0.4.
            shot_noise_std_beta (float): Shot noise beta. Default 0.0.
            wavelengths (list): Wavelengths.
            red_response (list): Red channel spectral response.
            green_response (list): Green channel spectral response.
            blue_response (list): Blue channel spectral response.
        """
        super().__init__(size=size, res=res)

        self.bit = bit
        self.nbit_max = 2**bit - 1
        self.black_level = black_level
        self.bayer_pattern = bayer_pattern

        # Noise parameters
        self.iso_base = iso_base
        self.readnoise_std = read_noise_std
        self.shotnoise_std_alpha = shot_noise_std_alpha
        self.shotnoise_std_beta = shot_noise_std_beta

        # Spectral response curves
        self.wavelengths = wavelengths
        if self.wavelengths is not None:
            green_sum = sum(green_response)
            self.red_response = torch.tensor(red_response) / green_sum
            self.green_response = torch.tensor(green_response) / green_sum
            self.blue_response = torch.tensor(blue_response) / green_sum

        # ISP
        self.isp = InvertibleISP(
            bit=bit,
            black_level=black_level,
            bayer_pattern=bayer_pattern,
            white_balance=white_balance,
            color_matrix=color_matrix,
            gamma_param=gamma_param,
        )

    @classmethod
    def from_config(cls, sensor_file):
        """Create an RGBSensor from a JSON config file.

        Args:
            sensor_file: Path to JSON sensor config file.

        Returns:
            RGBSensor instance.
        """
        with open(sensor_file, "r") as f:
            config = json.load(f)

        return cls(
            size=config["sensor_size"],
            res=config["sensor_res"],
            bit=config["bit"],
            black_level=config["black_level"],
            bayer_pattern=config.get("bayer_pattern", "rggb"),
            white_balance=config.get("white_balance_d50", (2.0, 1.0, 1.8)),
            color_matrix=config.get("color_matrix_d50", None),
            gamma_param=config.get("gamma_param", 2.2),
            iso_base=config.get("iso_base", 100),
            read_noise_std=config.get("read_noise_std", 0.5),
            shot_noise_std_alpha=config.get("shot_noise_std_alpha", 0.4),
            shot_noise_std_beta=config.get("shot_noise_std_beta", 0.0),
            wavelengths=config.get("wavelengths", None),
            red_response=config.get("red_spectral_response", None),
            green_response=config.get("green_spectral_response", None),
            blue_response=config.get("blue_spectral_response", None),
        )

    def to(self, device):
        super().to(device)
        if self.wavelengths is not None:
            self.red_response = self.red_response.to(device)
            self.green_response = self.green_response.to(device)
            self.blue_response = self.blue_response.to(device)
        return self

    def response_curve(self, img_spectral):
        """Apply response curve to the spectral image to get the raw image.

        Args:
            img_spectral: Spectral image, shape (B, C, H, W), range [0, 1]

        Returns:
            img_raw: Raw image, shape (B, 3, H, W), range [0, 1]

        Reference:
            [1] Spectral Sensitivity Estimation Without a Camera. ICCP 2023.
            [2] https://github.com/COLOR-Lab-Eilat/Spectral-sensitivity-estimation
        """
        if self.wavelengths is not None:
            img_raw = torch.zeros(
                (
                    img_spectral.shape[0],
                    3,
                    img_spectral.shape[2],
                    img_spectral.shape[3],
                ),
                device=img_spectral.device,
            )
            img_raw[:, 0, :, :] = (
                img_spectral * self.red_response.view(1, -1, 1, 1)
            ).sum(dim=1)
            img_raw[:, 1, :, :] = (
                img_spectral * self.green_response.view(1, -1, 1, 1)
            ).sum(dim=1)
            img_raw[:, 2, :, :] = (
                img_spectral * self.blue_response.view(1, -1, 1, 1)
            ).sum(dim=1)
        else:
            assert img_spectral.shape[1] == 3, (
                "No spectral response curves provided, input image must have 3 channels"
            )
            img_raw = img_spectral

        return img_raw

    def unprocess(self, image, in_type="rgb"):
        """Unprocess an image to unbalanced RAW RGB space.

        Args:
            image: Tensor of shape (B, 3, H, W), range [0, 1]
            in_type: Input image type, either "rgb" or "linear_rgb"

        Returns:
            image: Tensor of shape (B, 3, H, W), range [0, 1] in raw space
        """
        isp = self.isp

        # Inverse gamma correction
        if in_type == "linear_rgb":
            pass
        elif in_type == "rgb":
            image = isp.gamma.reverse(image)
        else:
            raise ValueError(f"Invalid input type: {in_type}")

        # Inverse color correction matrix
        image = isp.ccm.reverse(image)

        # Inverse auto white balance
        image = isp.awb.reverse(image)  # (B, 3, H, W), [0, 1]

        return image

    def linrgb2raw(self, img_linrgb):
        """Convert linear RGB image to raw Bayer space.

        Args:
            img_linrgb: Tensor of shape (B, 3, H, W), range [0, 1]

        Returns:
            bayer_nbit: Tensor of shape (B, 1, H, W), range [~black_level, 2**bit - 1]
        """
        black_level = self.black_level
        bit = self.bit

        bayer_float = self.isp.demosaic.reverse(img_linrgb)
        bayer_nbit = bayer_float * (2**bit - 1 - black_level) + black_level
        bayer_nbit = torch.round(bayer_nbit)
        return bayer_nbit

    def simu_noise(self, img_raw, iso):
        """Simulate sensor noise considering sensor quantization and noise model.

        Args:
            img_raw: N-bit clean image, (B, C, H, W), range [0, 2**bit - 1]
            iso: (B,), range [0, 800]

        Returns:
            img_raw_noise: N-bit noisy image, (B, C, H, W), range [0, 2**bit - 1]

        Reference:
            [1] "Unprocessing Images for Learned Raw Denoising."
            [2] https://www.photonstophotos.net/Charts/RN_ADU.htm
            [3] https://www.photonstophotos.net/Investigations/Measurement_and_Sample_Variation.htm
            [4] https://www.dpreview.com/forums/thread/4669806
        """
        nbit_max = self.nbit_max
        black_level = self.black_level
        device = img_raw.device

        # Calculate noise standard deviation
        shotnoise_std = torch.clamp(
            self.shotnoise_std_alpha * torch.sqrt(torch.clamp(img_raw - black_level, min=0.0))
            + self.shotnoise_std_beta,
            0.0,
        )
        if (iso > 800).any():
            raise ValueError(f"Currently noise model only works for low ISO <= 800, got {iso}")
        gain_analog = 1.0  # we only measured analog gain = 1.0
        gain_digit = (iso / self.iso_base).view(-1, 1, 1, 1)
        noise_std = torch.sqrt(
            shotnoise_std**2 * gain_digit * gain_analog
            + self.readnoise_std**2 * gain_digit**2
        )

        # Sample random noise
        noise_sample = (
            torch.normal(mean=0.0, std=1.0, size=img_raw.size(), device=device)
            * noise_std
        )
        img_raw_noisy = img_raw + noise_sample

        # Clip and quantize
        img_raw_noisy = torch.clip(img_raw_noisy, 0.0, nbit_max)
        img_raw_noisy = torch.round(img_raw_noisy)
        return img_raw_noisy

    def sample_augmentation(self):
        """Randomly sample a set of augmentation parameters for ISP modules. Used for data augmentation during training."""
        self.isp.gamma.sample_augmentation()
        self.isp.ccm.sample_augmentation()
        self.isp.awb.sample_augmentation()

    def reset_augmentation(self):
        """Reset parameters for ISP modules. Used for evaluation."""
        self.isp.gamma.reset_augmentation()
        self.isp.ccm.reset_augmentation()
        self.isp.awb.reset_augmentation()

    # ===============================
    # Packing and unpacking
    # ===============================
    def process2rgb(self, image, in_type="rggb"):
        """Process an image to a RGB image.

        Args:
            image: Tensor of shape (B, 3, H, W), range [0, 1]
            in_type: Input image type, either "rggb" or "bayer"

        Returns:
            image: Tensor of shape (B, 3, H, W), range [0, 1]
        """
        # Process to RGB
        if in_type == "rggb":
            image = self.isp(self.rggb2bayer(image))
        elif in_type == "bayer":
            image = self.isp(image)
        else:
            raise ValueError(f"Invalid input type: {in_type}")

        return image

    def bayer2rggb(self, bayer_nbit):
        """Convert RAW bayer image to RAW RGGB image.

        Args:
            bayer_nbit: Tensor of shape (B, 1, H, W), range [~black_level, 2**bit - 1]

        Returns:
            rggb: Tensor of shape (B, 3, H, W), range [0, 1]
        """
        black_level = self.black_level
        bit = self.bit

        if len(bayer_nbit.shape) == 2:
            bayer_nbit = bayer_nbit.unsqueeze(0).unsqueeze(0)
            single_image = True
        else:
            single_image = False

        B, _, H, W = bayer_nbit.shape
        bayer_rggb = torch.zeros(
            (B, 4, H // 2, W // 2), dtype=bayer_nbit.dtype, device=bayer_nbit.device
        )

        bayer_rggb[:, 0, :, :] = bayer_nbit[:, 0, 0:H:2, 0:W:2]
        bayer_rggb[:, 1, :, :] = bayer_nbit[:, 0, 0:H:2, 1:W:2]
        bayer_rggb[:, 2, :, :] = bayer_nbit[:, 0, 1:H:2, 0:W:2]
        bayer_rggb[:, 3, :, :] = bayer_nbit[:, 0, 1:H:2, 1:W:2]

        # Data range [black_level, 2**bit - 1] -> [0, 1]
        rggb = (bayer_rggb - black_level) / (2**bit - 1 - black_level)

        if single_image:
            rggb = rggb.squeeze(0)

        return rggb

    def rggb2bayer(self, rggb):
        """Convert RGGB image to RAW Bayer.

        Args:
            rggb: Tensor of shape [4, H/2, W/2] or [B, 4, H/2, W/2], range [0, 1]

        Returns:
            bayer: Tensor of shape [1, H, W] or [B, 1, H, W], range [~black_level, 2**bit - 1]
        """
        black_level = self.black_level
        bit = self.bit

        if len(rggb.shape) == 3:
            rggb = rggb.unsqueeze(0)
            single_image = True
        else:
            single_image = False

        B, _, H, W = rggb.shape
        bayer = torch.zeros((B, 1, H * 2, W * 2), dtype=rggb.dtype).to(rggb.device)

        bayer[:, 0, 0 : 2 * H : 2, 0 : 2 * W : 2] = rggb[:, 0, :, :]
        bayer[:, 0, 0 : 2 * H : 2, 1 : 2 * W : 2] = rggb[:, 1, :, :]
        bayer[:, 0, 1 : 2 * H : 2, 0 : 2 * W : 2] = rggb[:, 2, :, :]
        bayer[:, 0, 1 : 2 * H : 2, 1 : 2 * W : 2] = rggb[:, 3, :, :]

        # Data range [0, 1] -> [0, 2**bit-1]
        # bayer = torch.round(bayer * (2**bit - 1 - black_level) + black_level)
        bayer = bayer * (2**bit - 1 - black_level) + black_level

        if single_image:
            bayer = bayer.squeeze(0)

        return bayer
