import logging

from geogst.core.geology.profiles.profilers import *
from geogst.core.geology.profiles.profiletraces import *
from geogst.core.geometries.projections.crs import Crs
from geogst.core.geometries.projections.project import project
from geogst.core.utils.arrays import *
from geogst.core.geology.profiles.profilers import LineProfiler
from geogst.core.geology.profiles.profiletraces import ZTrace

logger = logging.getLogger(__name__)


class GeoProfile:
    """
    Class representing the topographic and geological elements
    representing a single geological profile.
    """

    def __init__(
        self,
        topo_profile: Union[None, ZTrace] = None,
        points_projections: Union[
         None,
         Dict[
             Category,
             Union[
                  Dict[Category, PointTrace],
                  Dict[Category, PlaneTrace],
                  Dict[Category, Fault]
             ]
         ]
        ] = None,
        lines_intersections: Union[
         None,
         List[
             List[
                 Tuple[
                     Category, Optional[List[UnionPtSegment2D]]]
             ]
         ]
        ] = None,
        lines_intersections_with_attitudes: Union[
         None,
         List[
             Dict[
                 Category,
                 List[PlaneTrace]
             ]
         ]
        ] = None,
        polygons_intersections: Union[
         None,
         List[
             List[
                 Tuple[
                     Category,
                     Optional[
                         List[UnionPtSegment2D]
                     ]
                 ]
             ]
         ]
        ] = None,
    ):

        if topo_profile:
            check_type(topo_profile, "Topographic profile", ZTrace)

        self._topo_profile = topo_profile
        self._points_projections = points_projections
        self._lines_intersections = lines_intersections
        self._lines_intersections_with_attitudes = lines_intersections_with_attitudes
        self._polygons_intersections = polygons_intersections

    @property
    def points_projections(self):
        return self._points_projections

    @points_projections.setter
    def points_projections(self, value):
        self._points_projections = value

    @property
    def lines_intersections(self):
        return self._lines_intersections

    @lines_intersections.setter
    def lines_intersections(self, value):
        self._lines_intersections = value

    @property
    def lines_intersections_with_attitudes(self):
        return self._lines_intersections_with_attitudes

    @lines_intersections_with_attitudes.setter
    def lines_intersections_with_attitudes(self, value):
        self._lines_intersections_with_attitudes = value

    @property
    def polygons_intersections(self):
        return self._polygons_intersections

    @polygons_intersections.setter
    def polygons_intersections(self, value):
        self._polygons_intersections = value

    def has_topography(self) -> bool:
        """
        Check if geoprofile has topography set.
        """

        return self._topo_profile is not None

    def clear_topo_profile(self):
        """

        :return:
        """

        self._topo_profile = None

    def s_min(self):
        """

        :return:
        """

        return self._topo_profile.s_min()

    def s_max(self):
        """

        :return:
        """

        return self._topo_profile.s_max()

    def z_min(self):
        """

        :return:
        """

        return self._topo_profile.z_min()

    def z_max(self):
        """

        :return:
        """

        return self._topo_profile.z_max()

    def length_2d(self) -> numbers.Real:
        """
        Returns the 2D length of the profile.

        :return: the 2D profile length.
        :rtype: numbers.Real.
        """

        return self._topo_profile.s_length()


class GeoProfiles:
    """
    Represents an ensemble of Geoprofile elements, stored as a list.
    It is related to parallel profiles products.
    """

    def __init__(self,
                 profilers: Optional[List[LineProfiler]] = None,
                 crs: Union[None, str, Crs] = None,
                 topo_profiles: Optional[List[ZTrace]] = None,
                 wkt_crs: Optional[str] = None,
                 ):
        """
        Instantiates a GeoProfiles ensemble, either:
        - from a list of LineProfiler plus their CRS (the profiles' topography
          is then sampled from each profiler's own line, e.g. after DEM draping); or
        - from an already-computed list of ZTrace plus a WKT CRS string (e.g.
          when re-loading previously exported/computed profiles for plotting,
          with no LineProfiler/line geometry available).

        :param profilers: the source line profilers (mutually exclusive with topo_profiles).
        :param crs: the CRS of the profilers (required together with 'profilers').
        :param topo_profiles: already-computed topographic profiles (mutually exclusive with profilers).
        :param wkt_crs: the WKT CRS of the topo_profiles (required together with 'topo_profiles').
        """

        if profilers is not None and topo_profiles is not None:
            raise Exception("Provide either 'profilers' or 'topo_profiles', not both")

        if profilers is not None:

            if crs is None:
                raise Exception("'crs' is required when 'profilers' is provided")

            self._profilers = Profilers(profilers)
            self._crs = crs if isinstance(crs, Crs) else Crs(crs)
            self._topo_profiles = []
            for ndx, line_profiler in enumerate(profilers):
                z_trace = ZTrace.from_mline(
                    mline=line_profiler.line,
                    wkt_crs=self._crs.to_wkt(),
                    rec_id=ndx,
                )
                self._topo_profiles.append(z_trace)

        elif topo_profiles is not None:

            crs_value = wkt_crs if wkt_crs is not None else crs
            if crs_value is None:
                raise Exception("'wkt_crs' is required when 'topo_profiles' is provided")

            self._profilers = Profilers([])
            self._crs = crs_value if isinstance(crs_value, Crs) else Crs(crs_value)
            self._topo_profiles = list(topo_profiles)

        else:
            raise Exception("Either 'profilers' or 'topo_profiles' must be provided")

        self._points_projections = {}
        self._points_meta = {}
        self._lines_intersections = None
        self._lines_with_attitudes_intersections = None
        self._polygons_intersections = None

    @property
    def profilers(self) -> Profilers:
        return self._profilers

    @property
    def wkt_crs(self) -> Optional[str]:
        return self._crs.to_wkt()

    @property
    def crs(self):
        return self._crs

    @property
    def topo_profiles(self):
        return self._topo_profiles

    @property
    def points_projections(self):
        return self._points_projections

    @property
    def points_meta(self):
        return self._points_meta

    @property
    def lines_intersections(self):
        return self._lines_intersections

    @property
    def lines_with_attitudes_intersections(self):
        return self._lines_with_attitudes_intersections

    @property
    def polygons_intersections(self):
        return self._polygons_intersections

    def sample_grid(self,
                                    grid: Grid,
                                    sampling_distance: Optional[numbers.Real] = None
                                    ) -> Error:
        """
        Set profile from a grid.

        :param grid: the source grid.
        :param sampling_distance: the grid sampling distance.
        :return: an error status.
        """

        try:

            result, err = self.profilers.sample_grid(
                grid=grid,
                sampling_distance=sampling_distance)

            if err:
                return err

            self._topo_profiles = result

            return Error()

        except Exception as e:

            return Error(
                True,
                caller_name(),
                e,
                traceback.format_exc())

    def intersect_lines(self,
        lines: Dict[Category, List[Ln]],
        lines_crs: Union[str, Crs],
        progress_callback: Optional[Callable[[numbers.Integral, numbers.Integral], None]] = None,
    ) -> Error:
        """
        Calculates the geoprofiles intersections with a set of lines.
        Note: the intersections are calculated in a 2D plane, not 3D.

        :param lines: the lines to intersect profile with
        :param progress_callback: optional callable invoked as
            progress_callback(done, total) after each of the 'total' parallel
            profiles is intersected, e.g. to drive a progress bar.
        :return: the error status.
        """

        try:

            # Projecting lines into geoprofiles CRS

            if not isinstance(lines_crs, Crs):
                lines_crs = Crs(lines_crs)

            geoprofiles_crs_projected_lines = {}
            for category, lines_geoms in lines.items():
                geoprofiles_crs_projected_lines_per_category = project(
                    lines_geoms,
                    lines_crs,
                    self.crs,
                )
                geoprofiles_crs_projected_lines[category] = geoprofiles_crs_projected_lines_per_category

            '''
            List[List[Dict[Category, List[UnionPtSegment2D]]]]
            first list is per geoprofile line
            second list is per geoprofile segment
            dictionary of categories versus list of points or segments
            '''

            geoprofiles_lines_intersections = self.profilers.intersect_lines(
                lines=geoprofiles_crs_projected_lines,
                progress_callback=progress_callback,
            )

            geoprofiles_valid_lines_intersections = self.profilers.parse_multipoints_intersections(
                geoprofiles_lines_intersections,
            )

            logger.debug("lines intersections:\n%s", geoprofiles_valid_lines_intersections)

            self._lines_intersections = geoprofiles_valid_lines_intersections

            return Error()

        except Exception as e:

            return Error(
                True,
                caller_name(),
                e,
                traceback.format_exc())

    def intersect_lines_with_attitudes(self,
        lines: Dict[Category, List[Tuple[Plane, List[Ln]]]]
    ) -> Error:
        """
        Calculates the profiles intersections with a set of lines.
        Note: the intersections are intended flat (in a 2D plane, not 3D).

        :param lines: the lines to intersect profile with
        :return: the error status.
        """

        try:

            intersections, err = self.profilers.intersect_lines_with_attitudes(
                lines=lines
            )

            if err:
                return err

            self._lines_with_attitudes_intersections = intersections

            return Error()

        except Exception as e:

            return Error(
                True,
                caller_name(),
                e,
                traceback.format_exc())

    def intersect_polygons(self,
        polygons: Dict[Category, List[Polygon]],
        polygons_crs: Union[str, Crs],
        progress_callback: Optional[Callable[[numbers.Integral, numbers.Integral], None]] = None,
    ) -> Error:
        """
        Calculates the profiles intersections with a set of polygon.
        Note: the intersections are intended flat (in a 2D plane, not 3D).

        :param polygons: the list of polygons to intersect with profile
        :param polygons_crs: the CRS of the polygons
        :param progress_callback: optional callable invoked as
            progress_callback(done, total) after each of the 'total' parallel
            profiles is intersected, e.g. to drive a progress bar.
        :return: the error status.
        """

        try:

            if not isinstance(polygons_crs, Crs):
                polygons_crs = Crs(polygons_crs)

            # Projecting polygons into geoprofiles CRS

            geoprofiles_crs_projected_polygons = {}
            for category, polygons_geoms in polygons.items():
                geoprofiles_crs_projected_polygons_per_category = project(
                    polygons_geoms,
                    polygons_crs,
                    self._crs,
                )
                geoprofiles_crs_projected_polygons[category] = geoprofiles_crs_projected_polygons_per_category

            geoprofiles_polygons_intersections = self.profilers.intersect_polygons(
                polygons=geoprofiles_crs_projected_polygons,
                progress_callback=progress_callback,
            )

            geoprofiles_valid_polygons_intersections = self.profilers.parse_multipoints_intersections(
                geoprofiles_polygons_intersections,
            )

            self._polygons_intersections = geoprofiles_valid_polygons_intersections

            return Error()

        except Exception as e:

            return Error(
                True,
                caller_name(),
                e,
                traceback.format_exc())

    def __len__(self) -> numbers.Integral:
        """
        Return the number of geoprofiles.
        """

        return self.num_profiles()

    def __getitem__(self,
        ndx: numbers.Integral
    ) -> GeoProfile:
        """
        Return the GeoProfile associated to the provided index.

        :param ndx: the index of the geoprofile to extract.
        :return: the geoprofile associated to the provided index.
        """

        # parse points projections: self._points_projections is a dict keyed by
        # dataset name (e.g. "points", "attitudes"), each value being a list with
        # one entry per profile, so it must be indexed per key, not as a whole.
        if not self._points_projections:
            profile_points_projections = None
        else:
            profile_points_projections = {}
            for point_dataset_key, point_dataset_list in self._points_projections.items():
                if point_dataset_list and ndx < len(point_dataset_list):
                    profile_points_projections[point_dataset_key] = point_dataset_list[ndx]

        return GeoProfile(
            topo_profile=self._topo_profiles[ndx] if self._topo_profiles and ndx < self.num_profiles() else None,
            points_projections=profile_points_projections,
            lines_intersections=self._lines_intersections[ndx] if self._lines_intersections and ndx < len(self._lines_intersections) else None,
            lines_intersections_with_attitudes=self._lines_with_attitudes_intersections[ndx] if self._lines_with_attitudes_intersections and ndx < len(self._lines_with_attitudes_intersections) else None,
            polygons_intersections=self._polygons_intersections[ndx] if self._polygons_intersections and ndx < len(self._polygons_intersections) else None
        )

    def __iter__(self):
        """
        Iter the Geoprofiles instance.
        """

        return (self[ndx] for ndx in range(len(self)))

    def have_topographies(self) -> bool:
        """
        Check if topographies defined.
        """

        return len(self._topo_profiles) > 0

    def num_profiles(self) -> numbers.Integral:
        """
        Returns the number of profiles in the geoprofile set.

        :return: number of profiles in the geoprofile set.
        :rtype: numbers.Integral.
        """

        return len(self.topo_profiles)

    def s_min(self) -> numbers.Real:
        """

        :return:
        """

        return np.nanmin([z_trace.s_min() for z_trace in self._topo_profiles])

    def s_max(self) -> numbers.Real:
        """

        :return:
        """

        return np.nanmax([z_trace.s_max() for z_trace in self._topo_profiles])

    def z_min(self):
        """

        :return:
        """

        return np.nanmin([z_trace.z_min() for z_trace in self._topo_profiles])

    def z_max(self):
        """

        :return:
        """

        return np.nanmax([z_trace.z_max() for z_trace in self._topo_profiles])

    def profiles_lengths_2d(self) -> List[numbers.Real]:
        """
        Returns the 2D lengths of the profiles.

        :return: the 2D profiles lengths.
        :rtype: list of numbers.Real values.
        """

        return [topoprofile.profile_length_2d() for topoprofile in self._topo_profiles]

    def profiles_lengths_3d(self) -> List[numbers.Real]:
        """
        Returns the 3D lengths of the profiles.

        :return: the 3D profiles lengths.
        :rtype: list of numbers.Real values.
        """

        return [topoprofile.profile_length_3d() for topoprofile in self._topo_profiles]

    def max_length_2d(self) -> Optional[numbers.Real]:
        """
        Returns the maximum 2D length of profiles.

        :return: the maximum profiles lengths.
        :rtype: an optional numbers.Real value.
        """

        lengths = self.profiles_lengths_2d()

        if lengths:
            return np.nanmax(lengths)
        else:
            return None


class MultiGridsProfiles:
    """
    Class storing multi grids profiles along a single spatial trace.
    """

    def __init__(self, profiler: LineProfiler):

        self._profiler = profiler

        self._zarray = ZArray()

    def sample_grid(self,
                    grid: Grid,
                    sampling_distance: Optional[numbers.Real] = None,
                    name: Optional[str] = ""
                    ) -> Error:
        """
        Set profile from a grid.

        :param grid: the source grid.
        :param sampling_distance: the grid sampling distance.
        :return: an error status.
        """

        try:

            result, err = self._profiler.sample_grid(
                grid=grid,
                sampling_distance=sampling_distance)

            if err:
                return err

            z_trace = result

            self._zarray.add_profile(
                z_trace=z_trace,
                name=name)

            return Error()

        except Exception as e:

            return Error(
                True,
                caller_name(),
                e,
                traceback.format_exc())

    def clear_profiles(self):
        """

        :return:
        """

        self._zarray = ZArray()

    def __getitem__(self,
        ndx: numbers.Integral
    ) -> GeoProfile:
        """
        Return the GeoProfile associated to the provided index.

        :param ndx: the index of the geoprofile to extract.
        :return: the geoprofile associated to the provided index.
        """

        return GeoProfile(
            topo_profile=self._zarray[ndx] if 0 <= ndx < self._zarray.num_profiles() else None,
        )

    def num_profiles(self) -> numbers.Integral:
        """
        Returns the number of profiles in the MultiGridsProfiles instance.

        :return: number of profiles in the MultiGridsProfiles instance.
        :rtype: numbers.Integral.
        """

        return 0 if self._zarray is None else self._zarray.num_profiles()

    def __len__(self) -> numbers.Integral:
        """
        Return the number of profiles in the MultiGridsProfiles instance.
        """

        return self.num_profiles()

    def __iter__(self):
        """
        Iter the Geoprofiles instance.
        """

        return (self[ndx] for ndx in range(len(self)))

    def s_min(self):
        """

        :return:
        """

        return self._zarray.s_min()

    def s_max(self):
        """

        :return:
        """

        return self._zarray.s_max()

    def z_min(self):
        """

        :return:
        """

        return self._zarray.z_min()

    def z_max(self):
        """

        :return:
        """

        return self._zarray.z_max()

    def length_2d(self) -> numbers.Real:
        """
        Returns the 2D length of the profile.

        :return: the 2D profile length.
        :rtype: numbers.Real.
        """

        return self.s_max()

