import json
from datetime import datetime
from typing import Any, List, Union

import numpy as np
import sqlalchemy as sa
from sqlalchemy import and_, or_
from sqlalchemy.orm import Session

from openlog.core import pint_utilities
from openlog.datamodel.assay.generic_assay import (
    AssayDatabaseDefinition,
    AssayDefinition,
    AssaySeriesType,
    DepthAssay,
    TimeAssay,
)


class SqlAlchemyAssayReader:
    """
    Helper class to read assay data from sqlalchemy session with a sqlalchemy table

    Args:
        hole_id: collar hole id for filter
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
        assay_definition: AssayDefinition used to describe assay contents
        assay_database_definition: AssayDatabaseDefinition  used to decribe how to insert assay into database
    """

    def __init__(
        self,
        hole_id: str,
        session: Session,
        assay_table: sa.Table,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
    ) -> None:
        self.hole_id = hole_id
        self._assay_table = assay_table
        self._session = session
        self._assay_definition = assay_definition
        self._assay_database_definition = assay_database_definition

    def get_distinct_col_values(
        self, column: str, related_column: str = ""
    ) -> np.array:
        """
        Returns distinct available value for assay column as numpy arrays

        Args:
            column (str): assay column
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row
        Returns:
            numpy.array: assay column value
        """
        filter_ = True
        filter_, query_col = self._add_y_column_filter_and_get_query_col(
            filter_, column
        )
        filter_ = self._add_remove_none_filter(filter_, column)
        filter_ = self._add_y_column_filter(filter_, related_column)

        all_values = (
            self._session.query(query_col)
            .filter(filter_)
            .distinct(query_col)
            .order_by(query_col)
            .all()
        )
        if len(all_values):
            full_array = np.array(all_values)
            y = np.array(full_array[:, 0])
            y = self._convert_from_base_unit_if_needed(column, y)
        else:
            y = np.array([])
        return y

    def _get_query_col(self, col_str: str) -> sa.column:
        """
        Get sqlalchemy column for a query for a specific column

        Returns sqlalchemy column defined in current sqlalchemy table or create a sqlalchemy column if not available

        Args:
            col_str: (str) column

        Returns: (sa.column) sqlalchemy column

        """

        if col_str in self._assay_table.c.keys():
            return self._assay_table.c[col_str]
        else:
            return sa.column(col_str)

    def _add_remove_none_filter(self, filter_: Any, column: str) -> Any:
        """
        Return a sqlalchemy statement with an added filter on None values

        Args:
            filter_: input sqlalchemy filter
            column: (str) column to filter for None values

        Returns: sqlalchemy statement with added filter

        """
        database_col = (
            self._assay_database_definition.y_col[column]
            if column in self._assay_database_definition.y_col
            else column
        )
        return and_(
            filter_,
            # WARNING : Need to use  `!=` and not `is not` or sqlachemy returns none values
            self._get_query_col(database_col) != None,
        )

    def get_full_table(self) -> list:
        """
        Get full assay table with all rows and all columns.
        """
        result = self._session.query(*self._assay_table.columns).all()
        return result

    def _add_y_column_filter_and_get_query_col(
        self, filter_: Any, column_: str
    ) -> (Any, Any):

        filter_ = self._add_y_column_filter(filter_, column_)

        if column_ in self._assay_database_definition.y_column_filter:
            query_col = sa.column(self._assay_database_definition.y_col[column_])
        else:
            database_col = (
                self._assay_database_definition.y_col[column_]
                if column_ in self._assay_database_definition.y_col
                else column_
            )
            query_col = self._get_query_col(database_col)
        return filter_, query_col

    def _add_y_column_filter(self, filter_: Any, column_: str) -> Any:
        if column_ in self._assay_database_definition.y_column_filter:
            row_filter = self._assay_database_definition.y_column_filter[column_]
            filter_ = and_(
                filter_,
                sa.column(row_filter) == column_,
            )
        return filter_

    def _convert_from_base_unit_if_needed(self, column: str, y: np.array) -> np.array:
        """
        Convert np.array from base unit if data is stored in base unit.

        Args:
            column: (str) column
            y: (np.array) input array

        Returns: (np.array) converted array

        """

        if column in self._assay_definition.columns:
            assay_column = self._assay_definition.columns[column]
            if (
                assay_column.unit
                and assay_column.stored_in_base_unit
                and assay_column.series_type != AssaySeriesType.POLAR
            ):
                y = pint_utilities.convert_from_base_unit(assay_column.unit, y)
        return y

    def extract_structural_values(self, serie: np.ndarray):
        """
        Convert spherical data (string) into an array of dict.
        """
        y_ = []
        for value in serie:
            if value == "None":
                y_.append(None)
                continue
            try:
                obj = json.loads(value)
                obj["dip"] = int(obj["dip"])
                obj["azimuth"] = float(obj["azimuth"])
                obj["polarity"] = int(obj["polarity"])

            except ValueError:
                azimuth, dip, polarity, type_ = (
                    value.replace("(", "").replace(")", "").split(",")
                )
                obj = {}
                obj["azimuth"] = float(azimuth)
                obj["dip"] = int(dip)
                obj["polarity"] = int(polarity)
                obj["type"] = type_.upper()

            y_.append(obj)

        return np.array(y_)

    @property
    def hole_col(self) -> sa.Column:
        return self._assay_table.c[self._assay_database_definition.hole_id_col]

    @property
    def x_col(self) -> sa.Column:
        return self._assay_table.c[self._assay_database_definition.x_col]

    @property
    def x_end_col(self) -> sa.Column:
        return self._assay_table.c[self._assay_database_definition.x_end_col]


class SqlAlchemyDiscreteAssayReader(SqlAlchemyAssayReader):
    """
    Helper class to read discrete assay data from sqlalchemy session with a sqlalchemy table

    Args:
        hole_id: collar hole id for filter
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
        assay_definition: AssayDefinition used to describe assay contents
        assay_database_definition: AssayDatabaseDefinition  used to describe how to insert assay into database
    """

    def __init__(
        self,
        hole_id: str,
        session: Session,
        assay_table: sa.Table,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
    ) -> None:
        super().__init__(
            hole_id, session, assay_table, assay_definition, assay_database_definition
        )

    def get_all_values(
        self,
        column: str,
        remove_none: bool = False,
        related_column: str = "",
        hole_filter: Any = True,
        return_hole_column: bool = False,
    ) -> List[np.ndarray]:
        """
        Returns all values available for assay column as numpy arrays

        Args:
            remove_none:
            column (str): assay column
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row
            hole_filter : (Any) True if hole_id is used for filtering. If False, no hole filter. If a list of str, filter on that list.
            return_hole_column: (bool) True if we want column describing hole, return 3 arrays instead of 2.
        """
        if hole_filter is True:
            filter_ = self.hole_col == self.hole_id
        elif isinstance(hole_filter, list):
            filter_ = self.hole_col.in_(hole_filter)
        else:
            filter_ = True

        if remove_none:
            filter_ = self._add_remove_none_filter(filter_, column)

        filter_, query_col = self._add_y_column_filter_and_get_query_col(
            filter_, column
        )

        filter_ = self._add_y_column_filter(filter_, related_column)
        all_values = (
            self._session.query(self.x_col, query_col)
            .filter(filter_)
            .order_by(self.x_col)
            .all()
        )
        x, y = self._convert_to_array(all_values)
        if (
            self._assay_definition.columns.get(column)
            and self._assay_definition.columns[column].series_type
            == AssaySeriesType.SPHERICAL
        ):
            y = self.extract_structural_values(y)

        y = self._convert_from_base_unit_if_needed(column, y)

        if return_hole_column:
            holes = (
                self._session.query(self.hole_col)
                .filter(filter_)
                .order_by(self.x_col)
                .all()
            )
            holes = np.array([hole[0] for hole in holes])
            return holes, x, y
        return x, y

    def _get_values_from_interval(
        self,
        column: str,
        x_min: Union[float, datetime],
        x_max: Union[float, datetime],
        remove_none: bool = False,
        related_column: str = "",
    ) -> (np.array, np.array):
        """
        Query session with sqlalchemy table class with x interval filter

        Args:
            column (str): assay column
            x_min: x minimum value
            x_max: x maximum value
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row

        Returns: tuple (x_value,y_values) as numpy array

        """
        filter_ = and_(
            x_min <= self.x_col,
            self.x_col <= x_max,
            self.hole_col == self.hole_id,
        )
        if remove_none:
            filter_ = self._add_remove_none_filter(filter_, column)

        filter_, query_col = self._add_y_column_filter_and_get_query_col(
            filter_, column
        )

        filter_ = self._add_y_column_filter(filter_, related_column)

        all_values = (
            self._session.query(self.x_col, query_col)
            .filter(filter_)
            .order_by(self.x_col)
            .all()
        )
        x, y = self._convert_to_array(all_values)
        y = self._convert_from_base_unit_if_needed(column, y)
        return x, y

    @staticmethod
    def _convert_to_array(values: List) -> (np.array, np.array):
        """
        Convert result of query to numpy array

        Args:
            values: query result

        Returns: tuple (x_value,y_values) as numpy array

        """
        if values:
            full_array = np.array(values)
            # Discrete assay : x has 1 dimension
            x = full_array[:, 0]
            # Define type from sqlalchemy result to avoid conversion of int into float
            y = np.array(full_array[:, 1], dtype=type(values[0][1]))
            return x, y
        else:
            return np.array([]), np.array([])


class SqlAlchemyExtendedAssayReader(SqlAlchemyAssayReader):
    """
    Helper class to read extended assay data from sqlalchemy session with a sqlalchemy table


    Args:
        hole_id: collar hole id for filter
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
        assay_definition: AssayDefinition used to describe assay contents
        assay_database_definition: AssayDatabaseDefinition  used to describe how to insert assay into database
    """

    def __init__(
        self,
        hole_id: str,
        session: Session,
        assay_table: sa.Table,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
    ) -> None:
        super().__init__(
            hole_id, session, assay_table, assay_definition, assay_database_definition
        )

    def get_all_values(
        self,
        column: str,
        remove_none: bool = False,
        related_column: str = "",
        hole_filter: Any = True,
        return_hole_column: bool = False,
    ) -> (np.array, np.array):
        """
        Returns all values available for assay column as numpy arrays

        Args:
            column (str): assay column
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row

        """
        if hole_filter is True:
            filter_ = self.hole_col == self.hole_id
        elif isinstance(hole_filter, list):
            filter_ = self.hole_col.in_(hole_filter)
        else:
            filter_ = True

        if remove_none:
            filter_ = self._add_remove_none_filter(filter_, column)

        filter_, query_col = self._add_y_column_filter_and_get_query_col(
            filter_, column
        )

        filter_ = self._add_y_column_filter(filter_, related_column)

        all_values = (
            self._session.query(
                self.x_col,
                self.x_end_col,
                query_col,
            )
            .filter(filter_)
            .order_by(self.x_col)
            .all()
        )
        x, y = self._convert_to_array(all_values)
        if (
            self._assay_definition.columns.get(column)
            and self._assay_definition.columns[column].series_type
            == AssaySeriesType.SPHERICAL
        ):
            y = self.extract_structural_values(y)
        y = self._convert_from_base_unit_if_needed(column, y)
        if return_hole_column:
            holes = (
                self._session.query(self.hole_col)
                .filter(filter_)
                .order_by(self.x_col)
                .all()
            )
            holes = np.array([hole[0] for hole in holes])
            return holes, x, y

        return x, y

    def _get_values_from_interval(
        self,
        column: str,
        x_min: Union[float, datetime],
        x_max: Union[float, datetime],
        remove_none: bool = False,
        related_column: str = "",
    ) -> (np.array, np.array):
        """
        Query session with sqlalchemy table class with x interval filter

        Args:
            column (str): assay column
            x_min: x minimum value
            x_max: x maximum value
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row

        Returns: tuple (x_value,y_values) as numpy array
        """
        filter_ = and_(
            self.hole_col == self.hole_id,
            or_(
                and_(x_min <= self.x_col, self.x_col <= x_max),
                and_(
                    x_min <= self.x_end_col,
                    self.x_end_col <= x_max,
                ),
                and_(x_min <= self.x_col, self.x_end_col <= x_max),
                and_(x_min >= self.x_col, self.x_end_col >= x_max),
            ),
        )
        if remove_none:
            filter_ = self._add_remove_none_filter(filter_, column)

        filter_, query_col = self._add_y_column_filter_and_get_query_col(
            filter_, column
        )

        filter_ = self._add_y_column_filter(filter_, related_column)

        all_values = (
            self._session.query(
                self.x_col,
                self.x_end_col,
                query_col,
            )
            .filter(filter_)
            .order_by(self.x_col)
            .all()
        )
        x, y = self._convert_to_array(all_values)
        y = self._convert_from_base_unit_if_needed(column, y)
        return x, y

    @staticmethod
    def _convert_to_array(values: List) -> (np.array, np.array):
        """
        Convert result of query to numpy array

        Args:
            values: query result

        Returns: tuple (x_value,y_values) as numpy array

        """
        if values:
            full_array = np.array(values)
            x = np.array(full_array[:, [0, 1]])
            # Define type from sqlalchemy result to avoid conversion of int into float
            y = np.array(full_array[:, 2], dtype=type(values[0][2]))
            return x, y
        else:
            return np.array([]), np.array([])


class SqlAlchemyDiscreteDepthAssay(SqlAlchemyDiscreteAssayReader, DepthAssay):
    """
    DepthAssay with discrete x values using a sqlalchemy connection to read data

    Inherits SqlAlchemyDiscreteAssayReader to read discrete data

    Args:
        hole_id: collar hole id for filter
        assay_definition: AssayDefinition
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
    """

    def __init__(
        self,
        hole_id: str,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
        session: Session,
        assay_table: sa.Table,
    ) -> None:
        DepthAssay.__init__(self, hole_id, assay_definition)
        SqlAlchemyDiscreteAssayReader.__init__(
            self,
            hole_id,
            session,
            assay_table,
            assay_definition,
            assay_database_definition,
        )

    def get_values_from_depth(
        self,
        column: str,
        depth_min: float,
        depth_max: float,
        remove_none: bool = False,
        related_column: str = "",
    ) -> (np.array, np.array):
        """
        Returns available values for assay column as numpy arrays from a depth interval

        Args:
            column (str): assay column
            depth_min (float): minimum depth
            depth_max (float): maximum depth
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row
        """
        return self._get_values_from_interval(
            column, depth_min, depth_max, remove_none, related_column
        )


class SqlAlchemyExtendedDepthAssay(SqlAlchemyExtendedAssayReader, DepthAssay):
    """
    DepthAssay with extended x values using a sqlalchemy connection to read data

    Inherits SqlAlchemyExtendedAssayReader to read extended data

    Args:
        hole_id: collar hole id for filter
        assay_definition: AssayDefinition
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
    """

    def __init__(
        self,
        hole_id: str,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
        session: Session,
        assay_table: sa.Table,
    ) -> None:
        DepthAssay.__init__(self, hole_id, assay_definition)
        SqlAlchemyExtendedAssayReader.__init__(
            self,
            hole_id,
            session,
            assay_table,
            assay_definition,
            assay_database_definition,
        )
        self._assay_table = assay_table
        self._session = session

    def get_values_from_depth(
        self,
        column: str,
        depth_min: float,
        depth_max: float,
        remove_none: bool = False,
        related_column: str = "",
    ) -> (np.array, np.array):
        """
        Returns available values for assay column as numpy arrays from a depth interval

        Args:
            column (str): assay column
            depth_min (float): minimum depth
            depth_max (float): maximum depth
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row
        """
        return self._get_values_from_interval(
            column, depth_min, depth_max, remove_none, related_column
        )


class SqlAlchemyDiscreteTimeAssay(SqlAlchemyDiscreteAssayReader, TimeAssay):
    """
    DepthAssay with discrete x values using a sqlalchemy connection to read data

    Inherits SqlAlchemyDiscreteAssayReader to read discrete data

    Args:
        hole_id: collar hole id for filter
        assay_definition: AssayDefinition
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
    """

    def __init__(
        self,
        hole_id: str,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
        session: Session,
        assay_table: sa.Table,
    ) -> None:
        TimeAssay.__init__(self, hole_id, assay_definition)
        SqlAlchemyDiscreteAssayReader.__init__(
            self,
            hole_id,
            session,
            assay_table,
            assay_definition,
            assay_database_definition,
        )
        self._assay_table = assay_table
        self._session = session

    def get_values_from_time(
        self,
        column: str,
        date_min: datetime,
        date_max: datetime,
        remove_none: bool = False,
        related_column: str = "",
    ) -> (np.array, np.array):
        """
        Returns available values for assay column as numpy arrays from a time interval.

        Args:
            column (str): assay column
            date_min (datetime): minimum datetime
            date_max (datetime): maximum datetime
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row

        """
        return self._get_values_from_interval(
            column, date_min, date_max, remove_none, related_column
        )


class SqlAlchemyExtendedTimeAssay(SqlAlchemyExtendedAssayReader, TimeAssay):
    """
    TimeAssay with extended x values using a sqlalchemy connection to read data

    Inherits SqlAlchemyExtendedAssayReader to read extended data

    Args:
        hole_id: collar hole id for filter
        assay_definition: AssayDefinition
        session: sqlalchemy session
        assay_table: sqlalchemy table class used for assay data request
    """

    def __init__(
        self,
        hole_id: str,
        assay_definition: AssayDefinition,
        assay_database_definition: AssayDatabaseDefinition,
        session: Session,
        assay_table: sa.Table,
    ):
        TimeAssay.__init__(self, hole_id, assay_definition)
        SqlAlchemyExtendedAssayReader.__init__(
            self,
            hole_id,
            session,
            assay_table,
            assay_definition,
            assay_database_definition,
        )
        self._assay_table = assay_table
        self._session = session

    def get_values_from_time(
        self,
        column: str,
        date_min: datetime,
        date_max: datetime,
        remove_none: bool = False,
        related_column: str = "",
    ) -> (np.array, np.array):
        """
        Returns available values for assay column as numpy arrays from a time interval.

        Args:
            column (str): assay column
            date_min (datetime): minimum datetime
            date_max (datetime): maximum datetime
            remove_none: (bool) remove none values (default False)
            related_column:  (str) related assay column (default ""). Used if an assay column is stored as row

        """
        return self._get_values_from_interval(
            column, date_min, date_max, remove_none, related_column
        )
