from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum

import numpy as np

from openlog.datamodel.assay.uncertainty import AssayColumnUncertainty


class AssayDomainType(Enum):
    TIME = "time"  # data is bound to time domain
    DEPTH = "depth"  # data is bound to depth domain
    # Data can't be represented by time and depth domain


class AssaySeriesType(Enum):
    NUMERICAL = "numerical"  # data represents a numerical value
    CATEGORICAL = "categorical"  # data is defined from a list of possible values
    NOMINAL = "nominal"  # data is defined by string value
    DATETIME = "datetime"  # data represents a datetime value
    IMAGERY = "imagery"  # data represents an imagery

    def python_type(self) -> type:
        if self == self.NUMERICAL:
            return float
        elif self == self.NOMINAL:
            return str
        elif self == self.CATEGORICAL:
            return str
        elif self == self.DATETIME:
            return datetime
        elif self == self.IMAGERY:
            return bytearray


class AssayDataExtent(Enum):
    DISCRETE = "discrete"  # data is bound to a single discrete value (evolution of value in time for example)
    EXTENDED = "extended"  # data is bound to 2 discrete value describing data extent (depth min and depth max for example)


@dataclass
class AssayColumn:
    """Class for assay column describing column name and type.

    Args:
        name (str): assay column name
        series_type (AssaySeriesType): assay series type (NUMERICAL/NOMINAL/CATEGORICAL/DATETIME)
        unit (str): assay unit
        category_name (str): category name for AssaySeriesType.CATEGORICAL
        stored_in_base_unit (bool): define if value is stored in base unit (default True)
        image_format_col (str): image format column for AssaySeriesType.IMAGERY (default "")
    """

    name: str
    series_type: AssaySeriesType
    unit: str = ""
    uncertainty: AssayColumnUncertainty = field(default_factory=AssayColumnUncertainty)
    category_name: str = ""
    stored_in_base_unit: bool = True
    image_format_col: str = ""


@dataclass
class AssayDefinition:
    """Class for assay definition describing nature of assay.

    Args:
        variable (str): assay variable
        display_name (str): assay display name
        domain (AssayDomainType): assay domain (TIME / DEPTH)
        data_extent (AssayDataExtent): assay data extent (DISCRETE/EXTENDED)
        columns {str: AssayColumn} : map by name of assay column
    """

    def __init__(
        self,
        variable: str,
        domain: AssayDomainType,
        data_extent: AssayDataExtent,
        columns: {str: AssayColumn},
        display_name: str = None,
    ) -> None:
        self.variable = variable
        self.domain = domain
        self.data_extent = data_extent
        self.columns = columns

        # Init display name with variable if not defined
        if display_name:
            self.display_name = display_name
        else:
            self.display_name = variable

    def get_uncertainty_columns(self) -> [str]:
        """
        Get unique uncertainty columns

        Returns: [str] unique uncertainty columns

        """
        result = []
        for col, column in self.columns.items():
            result += column.uncertainty.get_uncertainty_columns()
        return list(set(result))

    def get_image_format_columns(self) -> [str]:
        """
        Get unique image format columns

        Returns: [str] unique image format columns

        """
        result = []
        for _, column in self.columns.items():
            if column.image_format_col:
                result.append(column.image_format_col)
        return list(set(result))

    def get_plottable_columns(self) -> {str: AssayColumn}:
        """
        Get dict of assay column name and AssayColumn for plottable column

        Returns: [str:AssayColumn}

        """
        res = {}
        for col, column in self.columns.items():
            # No plot available for DATETIME
            valid = column.series_type != AssaySeriesType.DATETIME
            # Categorical and nominal data can only be displayed for extended extent
            if self.data_extent == AssayDataExtent.DISCRETE:
                valid &= column.series_type != AssaySeriesType.NOMINAL

            if valid:
                res[col] = column
        return res


@dataclass
class AssayDatabaseDefinition:
    """Class for assay database definition describing how assay is store in database.

    Args:
        table_name (str): assay table name
        hole_id_col (str): column for hole_id definition
        dataset_col (str): column for dataset definition
        x_col (str): column for x value definition (used for x_start value for AssayDataExtent.EXTENDED assays)
        y_col ({str: str}): column map for assay column value definition
        x_end_col (str) (optional): column for x_end value definition for AssayDataExtent.EXTENDED assays
        schema (str)(optional): schema name
    """

    def __init__(
        self,
        table_name: str,
        hole_id_col: str,
        dataset_col: str,
        x_col: str,
        y_col: {str: str},
        x_end_col: str = "",
        schema: str = "",
        y_column_filter: {str: str} = None,
    ) -> None:
        self.table_name = table_name
        self.hole_id_col = hole_id_col
        self.dataset_col = dataset_col
        self.x_col = x_col
        self.y_col = y_col

        if y_column_filter:
            self.y_column_filter = y_column_filter
        else:
            self.y_column_filter = {}

        self.x_end_col = x_end_col
        self.schema = schema

    def is_valid(self) -> bool:
        """
        Check if assay database definition is valid (no space in values and value defined)

        """
        result = True
        result &= " " not in self.table_name and self.table_name != ""
        result &= " " not in self.hole_id_col and self.hole_id_col != ""
        result &= " " not in self.dataset_col and self.dataset_col != ""
        result &= " " not in self.x_col and self.x_col != ""

        for col, col_table in self.y_col.items():
            result &= " " not in col_table and col_table != ""

        # optional values
        result &= " " not in self.x_end_col
        result &= " " not in self.schema
        return result


class GenericAssay:
    """
    Abstract interface for assay use.

    Methods to be implemented:

    - :meth:`get_all_values`: returns all available values for assay column as numpy arrays
    - :meth:`get_values_from_time`: returns available values for assay column as numpy arrays from a time interval
    - :meth:`get_values_from_depth`: returns available values for assay column as numpy arrays from a depth interval

    Attributes import_x_values / import_y_values are used for data importation
    import_y_value is a map of value for each available columns

    Args:
        hole_id: collar hole_id
        assay_definition: definition of assay (variable, domain, ...)
    """

    class InvalidInterface(Exception):
        pass

    class InvalidImportData(Exception):
        pass

    def __init__(self, hole_id: str, assay_definition: AssayDefinition):
        self.hole_id = hole_id
        self.assay_definition = assay_definition
        self.import_x_values = None
        self.import_y_values = {}

    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
        """
        raise GenericAssay.InvalidInterface()

    def get_all_values(
        self, column: str, remove_none: bool = False, related_column: str = ""
    ) -> (np.array, np.array):
        """
        Returns all values available for assay 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
        Returns:
            (numpy.array, numpy.array): assay value (x,y). x can be a array of tuple in case of extended data
        """
        raise GenericAssay.InvalidInterface()

    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 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
        Returns:
            (numpy.array, numpy.array): assay value (x,y). x can be a array of tuple in case of extended data
        """
        raise GenericAssay.InvalidInterface()

    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 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
        related_column: str = ""
        Returns:
            (numpy.array, numpy.array): assay value (x,y). x can be a array of tuple in case of extended data
        """
        raise GenericAssay.InvalidInterface()

    def check_import_data(self):
        """
        Check if import data is valid for current assay definition

        Raises GenericAssay.InvalidImportData otherwise

        """
        self._check_import_x_values()
        self._check_import_y_values()

    def _check_import_x_values(self):
        """
        Check if import x values are valid

        Raises GenericAssay.InvalidImportData otherwise

        """
        if type(self.import_x_values) is list:
            self.import_x_values = np.array(self.import_x_values)
        x_dimension = self.get_dimension(self.import_x_values.shape)
        # For discrete value, x represented only on dimension (time or depth)
        if self.assay_definition.data_extent == AssayDataExtent.DISCRETE:
            if x_dimension != 1:
                raise GenericAssay.InvalidImportData(
                    "Invalid discrete x import data, numpy array should have one "
                    "dimension"
                )

        # For extended value, x represented 2 dimensions (min/max)
        else:
            if x_dimension != 2:
                raise GenericAssay.InvalidImportData(
                    "Invalid extended x import data, numpy array should have 2 "
                    "dimensions"
                )

    def _check_import_y_values(self):
        """
        Check if import y values are valid

        Raises GenericAssay.InvalidImportData otherwise

        """
        expected_columns = self.assay_definition.columns.keys()
        nb_x = len(self.import_x_values)
        for col, vals in self.import_y_values.items():
            if (
                col not in expected_columns
                and col not in self.assay_definition.get_uncertainty_columns()
                and col not in self.assay_definition.get_image_format_columns()
            ):
                raise GenericAssay.InvalidImportData(
                    f"Invalid downhole data, column '{col}' not available in downhole data definition"
                )

            if type(vals) is list:
                vals = np.array(vals)

            if nb_x != len(vals):
                raise GenericAssay.InvalidImportData(
                    f"Invalid downhole data, x and y array for column '{col}' doesn't have same size"
                )

            y_dimension = self.get_dimension(vals.shape)
            if y_dimension != 1:
                raise GenericAssay.InvalidImportData(
                    f"Invalid y import data for column '{col}' , numpy array should have one dimension"
                )

    @staticmethod
    def get_dimension(shape) -> int:
        dimension = 1
        if len(shape) != 1:
            dimension = shape[1]
        return dimension


class TimeAssay(GenericAssay):
    """
    Abstract class for time assay

    Args:
        hole_id: collar hole_id
        assay_definition: definition of assay (variable, domain, ...)
    """

    def __init__(self, hole_id: str, assay_definition: AssayDefinition):
        super().__init__(hole_id, assay_definition)

    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 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
        Returns:
            (numpy.array, numpy.array): assay value (x,y). x can be a array of tuple in case of extended data
        """
        # Must be override in implementation class
        raise GenericAssay.InvalidInterface()


class DepthAssay(GenericAssay):
    """
    Abstract class for depth assay

    Args:
        hole_id: collar hole_id
        assay_definition: definition of assay (variable, domain, ...)
    """

    def __init__(self, hole_id: str, assay_definition: AssayDefinition):
        super().__init__(hole_id, assay_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 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
        Returns:
            (numpy.array, numpy.array): assay value (x,y). x can be a array of tuple in case of extended data
        """
        # Must be override in implementation class
        raise GenericAssay.InvalidInterface()
