from torch import Generator, Tensor

from smlmsim.camera.digitalization import (
    digitalization,
    digitalization_gain,
    digitalization_jac,
    reciprocal_digitalization_gain,
)
from smlmsim.camera.noise import (
    camera_noise,
    camera_noise_gain,
    camera_noise_jac,
    reciprocal_camera_noise_gain,
)


def apply_camera(
    y: Tensor,
    adu_baseline: float,
    camera_type: str,
    em_gain: float,
    inv_e_adu: float,
    quantum_efficiency: float,
    readout_noise: float,
    spurious_charge: float,
    gen: Generator,
):
    y = camera_noise(
        y,
        quantum_efficiency=quantum_efficiency,
        spurious_charge=spurious_charge,
        em_gain=em_gain,
        readout_noise=readout_noise,
        camera_type=camera_type,
        gen=gen,
    )
    y = digitalization(y, inv_e_adu=inv_e_adu, adu_baseline=adu_baseline)
    return y


def apply_camera_gain(
    y: Tensor,
    adu_baseline: float,
    camera_type: str,
    em_gain: float,
    inv_e_adu: float,
    quantum_efficiency: float,
):
    y = camera_noise_gain(
        y,
        quantum_efficiency=quantum_efficiency,
        em_gain=em_gain,
        camera_type=camera_type,
    )
    y = digitalization_gain(y, inv_e_adu=inv_e_adu, adu_baseline=adu_baseline)
    return y


def apply_camera_jac(
    y: Tensor,
    adu_baseline: float,
    camera_type: str,
    em_gain: float,
    inv_e_adu: float,
    quantum_efficiency: float,
):
    y = camera_noise_jac(
        y,
        quantum_efficiency=quantum_efficiency,
        em_gain=em_gain,
        camera_type=camera_type,
    )
    y = digitalization_jac(y, inv_e_adu=inv_e_adu, adu_baseline=adu_baseline)
    return y


def apply_reciprocal_camera_gain(
    y: Tensor,
    adu_baseline: float,
    camera_type: str,
    em_gain: float,
    inv_e_adu: float,
    quantum_efficiency: float,
):
    y = reciprocal_digitalization_gain(
        y, inv_e_adu=inv_e_adu, adu_baseline=adu_baseline
    )
    y = reciprocal_camera_noise_gain(
        y,
        quantum_efficiency=quantum_efficiency,
        em_gain=em_gain,
        camera_type=camera_type,
    )
    return y
