# Copyright (C) 2013-2023 Karlsruhe Institute of Technology
#
# This file is part of syris.
#
# This library is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library. If not, see <http://www.gnu.org/licenses/>.

"""Cameras used by experiments."""
import logging
import pkg_resources
import numpy as np
import quantities as q
import scipy.interpolate as interp
import syris.gpu.util as gutil
from syris import config as cfg
from syris.imageprocessing import bin_image, decimate
from syris.math import fwnm_to_sigma


LOG = logging.getLogger(__name__)


def is_fps_feasible(fps, exp_time):
    """Determine whether frame rate given by *fps* can be accomplished with
    the exposure time *exp_time*. It is only possible to set frame rates
    for which :math:`exposure\ time <= 1 / fps`.
    """
    return exp_time <= 1.0 / fps


class Camera(object):

    """Base class representing a camera."""

    def __init__(
        self,
        pixel_size,
        gain,
        dark_current,
        amplifier_sigma,
        bits_per_pixel,
        shape,
        quantum_efficiencies=None,
        wavelengths=None,
        exp_time=1 * q.s,
        fps=1 / q.s,
        dtype=np.ushort,
    ):
        """Create a camera with *pixel_size*, *gain* specifying :math:`\frac{counts}{e^-}`,
        *dark_current* as mean number of electrons present without incident light, *amplifier_sigma*
        is the sigma of the normally distributed electron noise given by camera electronics,
        *bit_per_pixel* is the number of bits holding the pixel grey value, *quantum_efficiencies*
        are the quantum efficiencies at *wavelengths*, *exp_time* is the exposure time and *fps* are
        Frames Per Second which are generated by the camera (exposure time is independent from fps,
        i.e.  e.g.  fps can be set to 1000 and exposure time to 1 \mu s, but it cannot exceed
        :math:`1/fps` s).  *dtype* is the sensor output data type. If the values given are
        incompatible, the frame rate is adjusted to the exposure time.
        """
        self.pixel_size = pixel_size.simplified
        self.gain = gain
        self.dark_current = dark_current
        self.amplifier_sigma = amplifier_sigma
        self.bpp = bits_per_pixel
        self._quantum_efficiencies = quantum_efficiencies
        self._wavelengths = wavelengths
        self.dtype = dtype
        self.shape = shape
        self._last_input_shape = None
        self._psf = None

        if self._quantum_efficiencies is not None and self._wavelengths is not None:
            self._qe_tck = interp.splrep(
                self._wavelengths.rescale(q.nm).magnitude, self._quantum_efficiencies
            )
        if not is_fps_feasible(fps, exp_time):
            fps = 1 / exp_time.simplified
        self._exp_time = exp_time
        self._fps = fps

    @property
    def wavelengths(self):
        return self._wavelengths

    @property
    def exp_time(self):
        return self._exp_time

    @exp_time.setter
    def exp_time(self, exp_time):
        if not is_fps_feasible(self.fps, exp_time):
            fmt = "Exposure time {} not possible for FPS {}, setting FPS to {}"
            LOG.debug(fmt.format(exp_time, self.fps, 1 / exp_time.simplified))
            self._fps = 1 / exp_time.simplified
        self._exp_time = exp_time.simplified

    @property
    def fps(self):
        return self._fps

    @fps.setter
    def fps(self, fps):
        if not is_fps_feasible(fps, self.exp_time):
            fmt = "FPS {} not possible for exposure time {}, setting exposure time to {}"
            LOG.debug(fmt.format(fps, self.exp_time, 1 / fps.simplified))
            self._exp_time = 1 / fps.simplified
        self._fps = fps.simplified

    @property
    def max_grey_value(self):
        return 2 ** self.bpp - 1

    def get_quantum_efficiency(self, wavelength):
        """Get quantum efficiency [dimensionless] at *wavelength*."""
        return interp.splev(wavelength.rescale(q.nm).magnitude, self._qe_tck)

    def get_image(self, photons, shot_noise=True, amplifier_noise=True, psf=True, queue=None):
        """Get digital counts image from incoming *photons*. The resulting image is based on the
        incoming photons and dark current. We apply noise based on EMVA 1288 standard according to
        which the variance :math:`\sigma_y^2 = K^2 ( \sigma_e^2 + \sigma_d^2 ) + \sigma_q^2`, where
        :math:`K` is the system gain, :math:`\sigma_e^2` is the poisson- distributed shot noise
        variance, :math:`\sigma_d^2` is the normal distributed electronics noise variance and
        :math:`\sigma_q^2` is the quantization noise variance. If *shot_noise* is False don't apply
        it. If *amplifier_noise* is False don't apply it as well. If *psf* is False don't apply the
        point spread function.
        """
        if self._last_input_shape != photons.shape:
            self._last_input_shape = photons.shape
            self._bin_factor = (photons.shape[0] / self.shape[0], photons.shape[1] / self.shape[1])

        if queue is None:
            queue = cfg.OPENCL.queue

        # Shot noise
        # Adjust dark current for later binning and gain
        dark = float(self.dark_current) / self._bin_factor[0] / self._bin_factor[1]
        electrons = dark + gutil.get_host(photons)

        if self._bin_factor != (1, 1):
            if psf:
                sigma = (fwnm_to_sigma(self._bin_factor[0]), fwnm_to_sigma(self._bin_factor[1]))
                small = decimate(electrons, self.shape, sigma=sigma, queue=queue)
            else:
                small = bin_image(electrons, self.shape, queue=queue)
            electrons = gutil.get_host(small)

        if shot_noise:
            electrons = np.random.poisson(electrons)

        if amplifier_noise and self.amplifier_sigma > 0:
            # Add electronics noise
            electrons = np.random.normal(electrons, self.amplifier_sigma)
            electrons[electrons < 0] = 0

        counts = self.gain * electrons

        # Cut the values beyond the maximum represented grey value given by
        # bytes per pixel.
        counts[counts > self.max_grey_value] = self.max_grey_value

        # Apply quantization noise
        return counts.astype(self.dtype)


def make_pco_dimax():
    """Make a pco.dimax camera."""
    lam, qe = np.load(
        pkg_resources.resource_filename(__name__, "data/dimax_quantum_efficiencies.npy")
    )
    lam = lam * q.m

    # Use a power of two padded value so that it's easier to use with FFT
    return Camera(
        11 * q.um, 0.1, 530.0, 23.0, 12, (2048, 2048), quantum_efficiencies=qe, wavelengths=lam
    )
