from datetime import datetime, timezone
from pathlib import Path
from typing import Generator

import numpy as np
try:
    import pandas as pd
except ImportError:
    from ..pymesh.stubs import pandas as pd

from .mesh_driver import MeshDriver, DatasetGroup
from ..._pytuflow_types import TimeLike

try:
    from qgis.core import (QgsMeshDatasetIndex, QgsMesh, QgsMeshSpatialIndex, QgsPointXY, QgsInterval,
                           QgsGeometry, QgsLineString, QgsApplication, QgsMeshLayer)
    from .scalar_mesh_result import ScalarMeshResult
    from .vector_mesh_result import VectorMeshResult
    from .mesh_result import MeshResult
    from .mesh_geom import mesh_intersects, IntersectResult

    has_qgis = True
except ImportError:
    has_qgis = False
    MeshResult = 'MeshResult'
    IntersectResult = 'IntersectResult'
    QgsMeshDatasetIndex = 'QgsMeshDatasetIndex'
    QgsMesh = 'QgsMesh'
    QgsPointXY = 'QgsPointXY'
    QgsInterval = 'QgsInterval'
    QgsGeometry = 'QgsGeometry'
    QgsLineString = 'QgsLineString'
    QgsApplication = 'QgsApplication'
    QgsMeshLayer = 'QgsMeshLayer'
    ScalarMeshResult = 'ScalarMeshResult'
    VectorMeshResult = 'VectorMeshResult'


Point = tuple[float, float]

from ...util import pytuflow_logging
logger = pytuflow_logging.get_logger()


class QgisMeshDriver(MeshDriver):
    DRIVER_SOURCE = 'qgis'

    def __init__(self, mesh: Path):
        super().__init__(mesh)
        self.lyr = None
        self.dp = None
        self.si = None
        self.qgsmesh = None
        self.reference_time = datetime(1990, 1, 1, tzinfo=timezone.utc)
        self.start_end_locs = []
        self._point_results = []
        self._linestrings = []
        self._line_results = []
        self._painter = None

    def __repr__(self):
        return f'<{self.__class__.__name__} {self.mesh.stem}>'

    def data_groups(self) -> Generator[DatasetGroup, None, None]:
        if not self.lyr:
            raise RuntimeError('Layer not loaded.')

        dataset_group_count = self.lyr.datasetGroupCount()
        i = -1
        while dataset_group_count:
            i += 1
            ind = QgsMeshDatasetIndex(i, 0)
            grp = self.lyr.datasetGroupMetadata(ind)
            name = grp.name()
            if not name:
                continue
            dataset_group_count -= 1
            type_ = 'vector' if grp.isVector() else 'scalar'
            times = [self.lyr.datasetMetadata(QgsMeshDatasetIndex(ind.group(), i)).time() for i in range(self.lyr.datasetCount(ind))]
            vert_lyr_count = grp.maximumVerticalLevelsCount()
            yield DatasetGroup(name, type_, times, vert_lyr_count)

    def init_mesh_layer(self, name: str):
        if not has_qgis:
            raise ImportError('QGIS python libraries are not installed or cannot be imported.')
        if not QgsApplication.instance():
            raise RuntimeError('QGIS application instance not found.')

        self.lyr = QgsMeshLayer(str(self.mesh), name, 'mdal')
        if not self.lyr.isValid():
            raise RuntimeError(f'Failed to load mesh layer {self.mesh}')

        self.dp = self.lyr.dataProvider()

    def load(self):
        if not has_qgis:
            raise ImportError('QGIS python libraries are not installed or cannot be imported.')

        target_dataset_count = self.lyr.datasetGroupCount()

        i = -1
        while target_dataset_count:
            i += 1
            ind = QgsMeshDatasetIndex(i, 0)
            grp = self.lyr.datasetGroupMetadata(ind)
            if not grp.name():
                continue
            target_dataset_count -= 1
            if grp.isTemporal():
                if grp.referenceTime().isValid():
                    self.reference_time = grp.referenceTime().toPyDateTime()
                    if self.reference_time.tzinfo is None:
                        self.reference_time = self.reference_time.replace(tzinfo=timezone.utc)
                    self.has_inherent_reference_time = True
                    break

    def init_spatial_index(self):
        if self.si:
            return

        if not self.lyr:
            raise RuntimeError('Layer not loaded.')

        if not self.dp:
            self.dp = self.lyr.dataProvider()

        self.qgsmesh = QgsMesh()
        self.dp.populateMesh(self.qgsmesh)
        self.si = QgsMeshSpatialIndex(self.qgsmesh)

    def group_index_from_name(self, data_type: str, **kwargs) -> int:
        from ..map_output import MapOutput  # import here to avoid circular import
        igrp = -1
        for i in range(self.lyr.datasetGroupCount()):
            ind = QgsMeshDatasetIndex(i, 0)
            ds_name = self.lyr.datasetGroupMetadata(ind).name()
            stnd_name = MapOutput._get_standard_data_type_name(ds_name)
            if stnd_name == data_type:
                igrp = i
                break

        if igrp == -1:
            raise ValueError(f'Dataset group not found for data type {data_type}')

        return igrp

    def maximum(self, data_type: str) -> float:
        idx = self.group_index_from_name(data_type)
        if idx == -1:
            raise ValueError(f'Data type {data_type} not found in mesh output {self.mesh.stem}')
        return self.lyr.datasetGroupMetadata(QgsMeshDatasetIndex(idx)).maximum()

    def minimum(self, data_type: str) -> float:
        idx = self.group_index_from_name(data_type)
        if idx == -1:
            raise ValueError(f'Data type {data_type} not found in mesh output {self.mesh.stem}')
        return self.lyr.datasetGroupMetadata(QgsMeshDatasetIndex(idx)).minimum()

    def data_point(self, point: Point, data_type: str, time: TimeLike, averaging_method: str | None = None, return_type: str = 'scalar') -> float | tuple[float, float]:
        self.init_spatial_index()
        igrp = self.group_index_from_name(data_type)

        if isinstance(time, datetime):
            reltime = (time - self.reference_time).total_seconds()
        else:
            reltime = time * 3600  # convert to seconds
        time_interval = QgsInterval(reltime)
        index = self.lyr.datasetIndexAtRelativeTime(time_interval, igrp)

        vector = self.lyr.datasetGroupMetadata(QgsMeshDatasetIndex(igrp)).isVector()
        if vector and return_type == 'vector':
            res = self._vector_point_result(QgsPointXY(*point))
        else:
            res = self._scalar_point_result(QgsPointXY(*point))

        value = res.value(index, averaging_method)
        if vector and isinstance(value, tuple) and return_type == 'scalar':
            value = np.sqrt(value[0] ** 2 + value[1] ** 2)
        return value

    def time_series(self, name: str, point: Point, data_type: str, averaging_method: str | None = None, return_type: str = 'scalar') -> pd.DataFrame:
        from ..map_output import MapOutput  # import here to avoid circular import
        self.init_spatial_index()
        igrp = self.group_index_from_name(data_type)
        stnd_name = MapOutput._get_standard_data_type_name(data_type)

        vector = self.lyr.datasetGroupMetadata(QgsMeshDatasetIndex(igrp)).isVector()
        if vector and return_type == 'vector':
            res = self._vector_point_result(QgsPointXY(*point))
        else:
            res = self._scalar_point_result(QgsPointXY(*point))

        data = []
        valid = False
        for i in range(self.dp.datasetCount(igrp)):
            index = QgsMeshDatasetIndex(igrp, i)
            ds = self.dp.datasetMetadata(index)
            value = res.value(index, averaging_method)
            if vector and isinstance(value, tuple) and return_type == 'scalar':
                value = np.sqrt(value[0] ** 2 + value[1] ** 2)
            time = ds.time()
            data.append((time, value))
            if not np.isnan(value):
                valid = True

        if not valid:
            return pd.DataFrame()

        df = pd.DataFrame(data, columns=['time', f'{name}/{stnd_name}'])
        df.set_index('time', inplace=True)
        return df

    def section(self, linestring: list[Point], data_type: str, time: TimeLike,
                averaging_method: str | None = None, return_type: str = 'scalar') -> pd.DataFrame:
        # noinspection DuplicatedCode
        df = pd.DataFrame()


        mesh_line = MeshLine(self, linestring, data_type, time, self.spherical, return_magnitude=(return_type == 'scalar'))

        # initialise start/end locations - for TUFLOW CATCH where results need to be stamped onto each other
        self.start_end_locs.clear()
        active = False
        start_loc, end_loc = None, None
        location = None

        # loop along line and get data
        data_ = []
        valid = False
        for location, mesh_result in mesh_line.results_along_line():
            # get value from mesh
            value = mesh_result.value(mesh_line.index, averaging_method)
            if isinstance(value, tuple):
                value = np.sqrt(value[0] ** 2 + value[1] ** 2)
            data_.append((location.dist1, value))
            if location.type == 'face':
                data_.append((location.dist2, value))

            # check if the returned data will have something valid in it
            if not np.isnan(value):
                valid = True

            # check if mesh is active so start / end locations can be recorded
            if mesh_result.active and not active:
                active = True
                start_loc = location.start_side
            elif not mesh_result.active and active:
                active = False
                end_loc = location.start_side
                self.start_end_locs.append((start_loc, end_loc))

        # check if mesh is active at the end of the line (close it off if it is still active)
        if active and location:  # location should be assigned if active is True
            end_loc = location.end_side
            self.start_end_locs.append((start_loc, end_loc))

        if valid:
            df = pd.DataFrame(np.array(data_), columns=['offset', data_type])
            df.set_index('offset', inplace=True)

        return df

    def curtain(self, linestring: list[Point], data_type: str, time: TimeLike) -> pd.DataFrame:
        # noinspection DuplicatedCode
        df = pd.DataFrame()
        mesh_line = MeshLine(self, linestring, data_type, time, self.spherical, vel_to_vec_vel=True)

        # initialise start/end locations - for TUFLOW CATCH where results need to be stamped onto each other
        self.start_end_locs.clear()
        active = False
        start_loc, end_loc = None, None
        location = None

        # loop through points and extract results
        data_ = []
        valid = False
        line_vector = []
        is_vector = False
        for i, (location, mesh_result) in enumerate(mesh_line.results_along_line()):
            if mesh_result.DATA_TYPE == 'Vector':
                is_vector = True
            if location.type == 'vertex' and (i == 0 or i == len(mesh_line.intersects.intersects)):  # ignore first and last points
                continue
            order_switch = False
            for z, value in mesh_result.vertical_values(mesh_line.index, 'stepped'):
                line_vector.append((location.intersect_vector_perp, location.intersect_vector))
                line_vector.append((location.intersect_vector_perp, location.intersect_vector))
                if order_switch:
                    data_.append((location.end_side, z, value))
                    data_.append((location.start_side, z, value))
                else:
                    data_.append((location.start_side, z, value))
                    data_.append((location.end_side, z, value))

                order_switch = not order_switch

                # check if the returned data will have something valid in it
                if isinstance(value, tuple):
                    value = value[0]
                if not np.isnan(value):
                    valid = True

            # check if mesh is active so start / end locations can be recorded
            if mesh_result.active and not active:
                active = True
                start_loc = location.start_side
            elif not mesh_result.active and active:
                active = False
                end_loc = location.start_side
                self.start_end_locs.append((start_loc, end_loc))

        # check if mesh is active at the end of the line (close it off if it is still active)
        if active and location:  # location should be assigned if active is True
            end_loc = location.end_side
            self.start_end_locs.append((start_loc, end_loc))

        if valid:
            df = pd.DataFrame(data_, columns=['x', 'y', data_type])
            if is_vector:
                line_vector = np.array(line_vector).reshape((-1, 4))
                df[f'{data_type}_local'] = VectorMeshResult.convert_vector_to_local(df[data_type].to_numpy(), line_vector)

        return df

    def profile(self, point: Point, data_type: str, time: TimeLike, interpolation: str, return_type: str = 'scalar') -> pd.DataFrame:
        self.init_spatial_index()

        # get dataset index based on the time
        igrp = self.group_index_from_name(data_type)
        if isinstance(time, datetime):
            reltime = (time - self.reference_time).total_seconds()
        else:
            reltime = time * 3600  # convert to seconds
        time_interval = QgsInterval(reltime)
        index = self.lyr.datasetIndexAtRelativeTime(time_interval, igrp)

        vector = self.lyr.datasetGroupMetadata(QgsMeshDatasetIndex(igrp)).isVector()
        if vector and return_type == 'vector':
            res = self._vector_point_result(QgsPointXY(*point))
        else:
            res = self._scalar_point_result(QgsPointXY(*point))

        data_ = []
        valid = False
        for z, value in res.vertical_values(index, 'stepped'):
            if isinstance(value, tuple):
                value = np.sqrt(value[0] ** 2 + value[1] ** 2)
            data_.append((value, z))
            if not np.isnan(value):
                valid = True

        df = pd.DataFrame()
        if valid:
            df = pd.DataFrame(data_, columns=[data_type, 'elevation'])
            df.set_index('elevation', inplace=True)

        return df

    def _scalar_point_result(self, point: QgsPointXY) -> ScalarMeshResult:
        res = ScalarMeshResult(self.lyr, self.qgsmesh, self.dp, self.si, QgsPointXY(*point))
        if res in self._point_results:
            res = self._point_results[self._point_results.index(res)]
        else:
            self._point_results.append(res)
        return res

    def _vector_point_result(self, point: QgsPointXY) -> ScalarMeshResult:
        res = VectorMeshResult(self.lyr, self.qgsmesh, self.dp, self.si, QgsPointXY(*point))
        if res in self._point_results:
            res = self._point_results[self._point_results.index(res)]
        else:
            self._point_results.append(res)
        return res


class MeshLine:
    """Class to represent extracting mesh results along a line-string. This class is used to perform
    common tasks for extracting mesh results along a line-string, such as finding the intersections and initialising
    the mesh result class.
    """

    def __init__(self, driver: QgisMeshDriver, linestring: list[Point], data_type: str, time: TimeLike, spherical: bool,
                 return_magnitude: bool = False, **kwargs):
        self.driver = driver  # parent class
        self.linestring = QgsLineString([QgsPointXY(*pnt) for pnt in linestring])
        self.data_type = data_type
        self.time = time
        self.spherical = spherical
        self.return_magnitude = return_magnitude

        self.driver.init_spatial_index()
        self.lyr = self.driver.lyr
        self.dp = self.driver.dp
        self.si = self.driver.si
        self.qgsmesh = self.driver.qgsmesh
        self.igrp = self.driver.group_index_from_name(data_type, **kwargs)

        # get dataset index based on the time
        if isinstance(time, datetime):
            reltime = (time - self.driver.reference_time).total_seconds()
        else:
            reltime = time * 3600  # convert to seconds
        time_interval = QgsInterval(reltime)
        self.index = self.lyr.datasetIndexAtRelativeTime(time_interval, self.igrp)
        linestring = QgsLineString([QgsPointXY(*pnt) for pnt in linestring])

        # get mesh intersections
        if linestring in self.driver._linestrings:
            self.intersects = self.driver._line_results[self.driver._linestrings.index(linestring)]
        else:
            self.intersects = mesh_intersects(self.qgsmesh, self.si, linestring, self.spherical)
            self.driver._linestrings.append(linestring)
            self.driver._line_results.append(self.intersects)

    def results_along_line(self) -> Generator[tuple[IntersectResult, MeshResult], None, None]:
        """Generator to yield the locations along the line-string."""
        grp = self.lyr.datasetGroupMetadata(self.index)
        for location in self.intersects.iter(self.dp.datasetGroupMetadata(self.igrp)):
            # mesh result class
            if grp.isVector() and not self.return_magnitude:
                mesh_result = VectorMeshResult(self.lyr, self.qgsmesh, self.dp, self.si, location.point)
            else:
                mesh_result = ScalarMeshResult(self.lyr, self.qgsmesh, self.dp, self.si, location.point)
            if mesh_result in self.driver._point_results:
                mesh_result = self.driver._point_results[self.driver._point_results.index(mesh_result)]
            else:
                self.driver._point_results.append(mesh_result)
            yield location, mesh_result
