import numpy as np

from openlog.datamodel.assay.generic_assay import AssayDataExtent


class GeoExtractor:
    """
    Abstract class designed to extract x,y,z geographical coordinates of an assay measure for a specific hole.
    Subclass should be used as an attribute of a GenericAssay instance and called during item construction (plot_item_factory.py).
    For depth domain only.
    """

    def __init__(self, assay) -> None:
        self.assay = assay
        self.geom = "proj_effective_geom"

    def get_trace_vertices(self, planned: bool = False) -> np.ndarray:
        """
        Return an array describing each vertex coordinates.
        """

        return NotImplementedError

    def get_planned_eoh(self):
        """
        Return altitude of planned EOH.
        """
        return NotImplementedError

    def get_altitude(self, planned: bool = False, remove_none: str = ""):
        """
        Args:
            - planned (bool) : True for planned trace
            - remove_none (str) : name of the column for filtering Nones, if empty no filter.
        """
        self.geom = "proj_planned_geom" if planned else "proj_effective_geom"
        if self.assay.assay_definition.data_extent == AssayDataExtent.DISCRETE:
            return self.get_discrete_altitude(remove_none)
        else:
            return self.get_extended_altitude(remove_none)

    def get_discrete_altitude(self, remove_none: str = "") -> np.ndarray:
        coords = self.get_discrete_coordinates(remove_none)
        if coords is not None:
            z = coords[:, 2]
            return z.astype(float)
        else:
            return

    def get_extended_altitude(self, remove_none: str = ""):
        coords = self.get_extended_coordinates(remove_none)
        if coords is not None:
            z = [np.array(coords)[:, i].astype(float) for i in [2, 5]]
            return z
        else:
            return

    def get_coordinates(self, planned: bool = False, remove_none: str = ""):
        self.geom = "proj_planned_geom" if planned else "proj_effective_geom"
        if self.assay.assay_definition.data_extent == AssayDataExtent.DISCRETE:
            return self.get_discrete_coordinates(remove_none)
        else:
            return self.get_extended_coordinates(remove_none)

    def get_discrete_coordinates(self, remove_none: str = ""):
        return NotImplementedError

    def get_extended_coordinates(self, remove_none: str = ""):
        return NotImplementedError


class SpatialiteGeoExtractor(GeoExtractor):
    def __init__(self, assay) -> None:
        super().__init__(assay)

    def get_trace_vertices(self, planned: bool = False):
        """
        Return an array describing each vertex coordinates in original SRID.
        Lines are from first to last vertex, columns ara x,y,z, cumulated length, assay length.
        """
        geom = "proj_planned_geom" if planned else "proj_effective_geom"

        q = f"""
            SELECT astext({geom}), eoh, planned_eoh from collar where hole_id = '{self.assay.hole_id}'
            """

        result = self.assay._session.execute(q).fetchall()[0]
        res = result[0]
        eoh = result[1]
        planned_eoh = result[2]
        if eoh is not None or planned_eoh is not None:
            eoh = max([val for val in [eoh, planned_eoh] if val is not None])

        if res is None:
            return

        # convert wkt to array : x,y,z,d
        res = res.split("(")[1]
        res = res.replace(")", "")
        res = res.split(", ")
        res = [txt.split(" ") for txt in res]
        # x,y,z
        arr = np.array(res).astype(float)
        # compute cumulated distance
        distance = 0
        dist_list = [0]
        for i in range(arr.shape[0] - 1):
            x1, y1, z1 = arr[i, :]
            x2, y2, z2 = arr[i + 1, :]
            d = np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2)
            distance += d
            dist_list.append(distance)

        # length correction (due to desurveying)
        coef = distance / eoh
        dist_list = np.array(dist_list)
        assay_dist = dist_list / coef
        arr = np.hstack(
            [arr, np.expand_dims(dist_list, axis=1), np.expand_dims(assay_dist, axis=1)]
        )

        return arr

    def get_planned_eoh(self):
        q = f"""
            SELECT st_z(geomfromtext(st_3dlineinterpolatepoint(proj_planned_geom, planned_eoh/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))), CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS c_z
            FROM collar
            WHERE hole_id = '{self.assay.hole_id}'
            """
        alt = self.assay._session.execute(q).fetchall()[0][0]
        return alt

    def get_discrete_coordinates(self, remove_none: str = "") -> np.ndarray:

        none_filter = ""
        if remove_none != "":
            none_filter = f"AND {remove_none} IS NOT NULL"
        q = f"""SELECT c_x, c_y, c_z FROM
                (SELECT {self.assay.hole_col.name},o.{self.assay.x_col.name}, max(ifnull(eoh, 0), ifnull(planned_eoh, 0)),
                st_x(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS c_x,
                st_y(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS c_y,
                st_z(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS c_z
                FROM (SELECT * FROM {self.assay._assay_table.name} where {self.assay.hole_col.name} = '{self.assay.hole_id}' {none_filter}) AS o
                LEFT JOIN collar AS c
                ON o.{self.assay.hole_col.name} = c.hole_id
                ORDER BY o.{self.assay.x_col.name} ) AS tmp

                """

        alt = self.assay._session.execute(q).fetchall()
        if len(alt) == 0:
            return None
        alt = np.array(alt)
        # if geometry doesn't exist
        if alt[0, 0] is None:
            return None
        alt = alt.astype(float)
        return alt

    def get_extended_coordinates(self, remove_none: str = ""):

        none_filter = ""
        if remove_none != "":
            none_filter = f"AND {remove_none} IS NOT NULL"

        q = f"""SELECT f_x, f_y, f_z, t_x, t_y, t_z FROM
                (SELECT {self.assay.hole_col.name},o.{self.assay.x_col.name}, o.{self.assay.x_end_col.name}, max(ifnull(eoh, 0), ifnull(planned_eoh, 0)),
                st_x(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS f_x,
                st_y(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS f_y,
                st_z(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS f_z,
                st_x(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_end_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS t_x,
                st_y(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_end_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS t_y,
                st_z(geomfromtext(st_3dlineinterpolatepoint({self.geom}, o.{self.assay.x_end_col.name}/max(ifnull(eoh, 0), ifnull(planned_eoh, 0))),  CASE WHEN SridIsProjected(srid) == 1 THEN srid ELSE 3857 END)) AS t_z
                FROM (SELECT * FROM {self.assay._assay_table.name} where {self.assay.hole_col.name} = '{self.assay.hole_id}' {none_filter}) AS o
                LEFT JOIN collar AS c
                ON o.{self.assay.hole_col.name} = c.hole_id
                ORDER BY o.{self.assay.x_col.name} ) AS tmp"""

        alt = self.assay._session.execute(q).fetchall()
        if len(alt) == 0:
            return None
        alt = np.array(alt)
        # if geometry doesn't exist
        if alt[0, 0] is None:
            return None
        alt = alt.astype(float)
        return alt


class XplordbGeoExtractor(GeoExtractor):
    def __init__(self, assay) -> None:
        super().__init__(assay)

    def get_trace_vertices(self, planned: bool = False):
        """
        Return an array describing each vertex coordinates.
        Lines are from first to last vertex, columns ara x,y,z, cumulated length, assay length.
        """
        geom = "proj_planned_geom" if planned else "proj_effective_geom"

        q = f"""
            select st_x(geom) as x, st_y(geom) as y, st_z(geom) as z, eoh
            from (select geom(st_dumppoints({geom})), greatest(d.eoh, d.planned_eoh) as eoh from display.display_collar as d
            join dh.collar as c on d.hole_id = c.hole_id
            where d.hole_id = '{self.assay.hole_id}') tmp

            """

        result = self.assay._session.execute(q).fetchall()
        if len(result) == 0:
            return

        arr = np.array(result).astype(float)
        eoh = float(arr[0, 3])
        arr = arr[:, :3]

        # compute cumulated distance
        distance = 0
        dist_list = [0]
        for i in range(arr.shape[0] - 1):
            x1, y1, z1 = arr[i, :]
            x2, y2, z2 = arr[i + 1, :]
            d = np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2)
            distance += d
            dist_list.append(distance)

        # length correction (due to desurveying)
        coef = distance / eoh
        dist_list = np.array(dist_list)
        assay_dist = dist_list / coef
        arr = np.hstack(
            [arr, np.expand_dims(dist_list, axis=1), np.expand_dims(assay_dist, axis=1)]
        )

        return arr

    def get_planned_eoh(self):
        q = f"""
            SELECT st_z(st_transform(ST_3DLineInterpolatePoint(ST_CurveToLine(proj_planned_geom), d.planned_eoh/greatest(d.eoh, d.planned_eoh)), st_srid(d.planned_geom))) AS c_z
            FROM display.display_collar as d
            JOIN dh.collar as c on d.hole_id = c.hole_id
            WHERE d.hole_id = '{self.assay.hole_id}'
            """
        alt = self.assay._session.execute(q).fetchall()[0][0]
        return alt

    def get_discrete_coordinates(self, remove_none: str = "") -> np.ndarray:

        none_filter = ""
        if remove_none != "":
            none_filter = f"AND {remove_none} IS NOT NULL"

        q = f"""SELECT c_x, c_y, c_z FROM
                (SELECT {self.assay.hole_col.name},o.{self.assay.x_col.name}, d.eoh,
                st_x(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS c_x,
                st_y(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS c_y,
                st_z(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS c_z
                FROM (SELECT * FROM assay.{self.assay._assay_table.name} where {self.assay.hole_col.name} = '{self.assay.hole_id}' {none_filter}) AS o
                LEFT JOIN display.display_collar AS d
                JOIN dh.collar as c on d.hole_id = c.hole_id
                ON o.{self.assay.hole_col.name} = d.hole_id
                ORDER BY o.{self.assay.x_col.name} ) AS tmp"""

        alt = self.assay._session.execute(q).fetchall()
        if len(alt) == 0:
            return None
        alt = np.array(alt)
        # if geometry doesn't exist
        if alt[0, 0] is None:
            return None
        alt = alt.astype(float)
        return alt

    def get_extended_coordinates(self, remove_none: str = ""):

        none_filter = ""
        if remove_none != "":
            none_filter = f"AND {remove_none} IS NOT NULL"

        q = f"""SELECT f_x, f_y, f_z, t_x, t_y, t_z FROM
                (SELECT {self.assay.hole_col.name},o.{self.assay.x_col.name}, o.{self.assay.x_end_col.name}, d.eoh,
                st_x(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS f_x,
                st_y(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS f_y,
                st_z(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS f_z,
                st_x(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_end_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS t_x,
                st_y(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_end_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS t_y,
                st_z(ST_3DLineInterpolatePoint(ST_CurveToLine({self.geom}), least(o.{self.assay.x_end_col.name}/greatest(d.eoh, d.planned_eoh), 1.0))) AS t_z
                FROM (SELECT * FROM assay.{self.assay._assay_table.name} where {self.assay.hole_col.name} = '{self.assay.hole_id}' {none_filter}) AS o
                LEFT JOIN display.display_collar AS d
                JOIN dh.collar as c on d.hole_id = c.hole_id
                ON o.{self.assay.hole_col.name} = d.hole_id
                ORDER BY o.{self.assay.x_col.name} ) AS tmp"""

        alt = self.assay._session.execute(q).fetchall()
        if len(alt) == 0:
            return None
        alt = np.array(alt)
        # if geometry doesn't exist
        if alt[0, 0] is None:
            return None
        alt = alt.astype(float)
        return alt
