# Copyright (C) 2024- Davide Mollica <davide.mollica@inaf.it>
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This file is part of iactsim.
#
# iactsim is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# iactsim 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with iactsim.  If not, see <https://www.gnu.org/licenses/>.

from pathlib import Path as _Path
from typing import (
    Tuple as _Tuple,
    List as _List,
    Union as _Union,
    Dict as _Dict,
    Any as _Any
)
import yaml as _yaml

import cupy as _cp
import numpy as _np

from ...io._path_loader import PathLoader as _PathLoader
from ...optics import AtmosphericTransmission as _AtmosphericTransmission

from ._os import AstriOpticalSystem as _AstriOpticalSystem
from ._camera import AstriCherenkovCamera as _AstriCherenkovCamera
from ..._iact import IACT as _IACT

class AstriTelescope(_IACT):
    """ASTRI telescope.

    Warning
    -------
    In order to perform camera simulation the `optical_system.build_camera_modules()` must be called before `cuda_init()` call.
    Do not change the number of pixels or the number of modules.

    See Also
    --------
        :py:class:`iactsim.IACT`:
    """
    def __init__(
        self,
        optical_system = None,
        camera: _AstriCherenkovCamera = None,
        position: _Tuple[float, float, float] | _List[float] | _np.ndarray = (0.,0.,0.),
        pointing: _Tuple[float, float] | _List[float] | _np.ndarray = (0,0),
    ):
        if optical_system is None:
            optical_system = _AstriOpticalSystem()
        
        if camera is None:
            camera = _AstriCherenkovCamera()

        super().__init__(
            optical_system = optical_system,
        )

        self.camera = camera

        # The telescope reference frame is defined such that when the telescope is
        # pointing at the horizon towards the North (i.e. alt: 0, az: 0), the axes are aligned as follows:
        # z-axis: aligned with the telescope optical axis and pointing North.
        # x-axis: perpendicular to the optical axis, lying in the horizontal plane and pointing West.
        # y-axis: perpendicular to the optical axis, pointing upwards.
        self.custom_rotation = _np.array([
            [ 0, 1, 0], # new X is old Y
            [-1, 0, 0], # new Y is old -X
            [ 0, 0, 1]  # new Z is old Z
        ])

        self.position = position
        self.pointing = pointing
    
    def update_camera_ucells(self):
        """Helper method to update the number of microcells of each SiPM following the definition in the optical system.

        Warning
        -------
        Changing the number of pixel or the number of modules is not supported.
        """
        n_ucells = []
        for s in self.optical_system:
            if s.name.startswith('PDM'):
                n_ucells.append(s._n_ucells_per_side**2)
        self.camera.n_ucells = _cp.repeat(_cp.asarray(n_ucells, dtype=_cp.int32), 64)

    def load_atmospheric_transmission(self, filepath: str):
        try:
            atm = _AtmosphericTransmission.load_simtel_file(filepath)
        except Exception as e:
            raise(RuntimeError("Unable to load atmospheric transmission from simtel atmospheric transimission file."))

        self.atmospheric_transmission = atm

    def configure(self, config: _Union[str, _Path, _Dict[str, _Any]]) -> None:
        """Configure the telescope with a yaml configuration file or dictionaries.
        The configuration files can contain more than a document to initialize also
        the optical system and the Cherenkov camera::

            ----
            # Telescope configuration
            # ...
            # ...
            ----
            # Optical system configuration
            # ...
            # ...
            ----
            # Camera configuration
            # ...
            # ...

        Parameters
        ----------
        config : str, path-like, or dict
            Path to the configuration file (str or pathlib.Path),
            or a dictionary containing the configuration.
        """
        configs = [None]*3

        if isinstance(config, (str, _Path)):
            with open(config, 'r') as file:
                tel_conf = _yaml.load_all(file, Loader=_PathLoader)
                for i,d in enumerate(tel_conf):
                    configs[i] = d
        elif isinstance(config, list):
            for i,d in enumerate(config):
                configs[i] = d
        elif isinstance(config, dict):
            configs[0] = config
        else:
            raise TypeError(
                "config must be a string (path), pathlib.Path, a list of dictionaries or a dictionary, "
                f"not {type(config)}"
            )

        n_documents = sum(d is not None for d in configs)

        if n_documents > 3:
            raise(RuntimeError('The configuration file can contain at most 3 document in the following order: telescope, optical system and camera.'))

        if n_documents == 0:
            raise(RuntimeError('The configuration file is empty.'))
        
        self._configure(configs[0])

        if configs[1] is not None:
            self.optical_system.configure(configs[1])
        
        if configs[2] is not None:
            self.camera.configure(configs[2])
        
        self.cuda_init()

    def _configure(self, tel_conf: _Dict[str, _Any]) -> None:
        """Configure the telescope with a dictionary.
        """
        # Assign only these
        assign_directly = [
            'altitude',
            'azimuth',
            'show_progress'
        ]

        for name in tel_conf.keys():
            value = tel_conf[name]
            if value is not None:
                if name in assign_directly:
                    setattr(self, name, value)
                if name == 'atmospheric_transmission':
                    self.load_atmospheric_transmission(value)