from dataclasses import dataclass
from enum import Enum

import numpy as np
import pandas as pd
from numpy import ma
from pandas.core.dtypes.common import is_numeric_dtype
from pandas.core.groupby import DataFrameGroupBy


class GapDataResolution(Enum):
    ACCEPT = "accept"
    REJECT = "reject"
    FORWARD_EXPANSION = "forward expansion"
    BACKWARD_EXPANSION = "backward expansion"
    NEAREST_NEIGHBOR = "nearest neighbor"


class OverlapDataResolution(Enum):
    REJECT = "reject"
    FORWARD_EXPANSION = "forward expansion"
    BACKWARD_EXPANSION = "backward expansion"
    NEAREST_NEIGHBOR = "nearest neighbor"


@dataclass
class ExtendedAssayColumn:
    hole_id_col: str
    from_col: str
    to_col: str
    val_cols: [str]
    update_status_col: str = ""


class ExtendedAssayInterpolation:
    class InvalidColumnData(Exception):
        pass

    def __init__(
        self,
        gap_resolution: GapDataResolution,
        overlap_resolution: OverlapDataResolution,
        columns: ExtendedAssayColumn,
    ) -> None:
        """
        Interpolator for extended assay data.

        Interpolate gap and overlap data :
        - accept
        - reject
        - forward expansion
        - backward expansion
        - nearest neighbor


        Args:
            gap_resolution: gap data resolution
            overlap_resolution: overlap data resolution
            columns: input dataframe columns
        """
        self._gap_resolution = gap_resolution
        self._overlap_resolution = overlap_resolution
        self._cols = columns

        # Define default value for update col
        if not self._cols.update_status_col:
            self._cols.update_status_col = "update_status"

    def interpolated_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Interpolate dataframe with current gap and overlap resolution option

        Args:
            df: input dataframe

        Returns: interpolated dataframe

        """
        # Check from and to col must have same dtype
        if not is_numeric_dtype(df.dtypes[self._cols.from_col]) or not is_numeric_dtype(
            df.dtypes[self._cols.to_col]
        ):
            raise ExtendedAssayInterpolation.InvalidColumnData(
                "From and to columns must have numeric dtype"
            )

        # Add update status col if not available
        if self._cols.update_status_col not in df.columns:
            df[self._cols.update_status_col] = None

        # Sort values
        df = df.sort_values(
            by=[self._cols.hole_id_col, self._cols.from_col, self._cols.to_col]
        )
        # group by hole id to check for gap and overlap
        for hole_id, group in df.groupby(self._cols.hole_id_col):
            # Overlap resolution must be done first (no data inserted, only row change)
            if self._overlap_resolution == OverlapDataResolution.BACKWARD_EXPANSION:
                df = self._expansion_overlap(df, group, True)
            elif self._overlap_resolution == OverlapDataResolution.FORWARD_EXPANSION:
                df = self._expansion_overlap(df, group, False)
            elif self._overlap_resolution == OverlapDataResolution.NEAREST_NEIGHBOR:
                df = self._nearest_neighbor_overlap(df, group)
            elif self._overlap_resolution == OverlapDataResolution.REJECT:
                df = self._rejection_overlap(df, group, hole_id)

            # Gap resolution done second because of data insert
            if self._gap_resolution == GapDataResolution.NEAREST_NEIGHBOR:
                df = self._nearest_neighbor_gap(df, group, hole_id)
            elif self._gap_resolution == GapDataResolution.BACKWARD_EXPANSION:
                df = self._expansion_gap(df, group, hole_id, True)
            elif self._gap_resolution == GapDataResolution.FORWARD_EXPANSION:
                df = self._expansion_gap(df, group, hole_id, False)
            elif self._gap_resolution == GapDataResolution.REJECT:
                df = self._rejection_gap(df, group, hole_id)

        # Sort data for new inserted rows
        return df.sort_values([self._cols.hole_id_col, self._cols.from_col])

    def _gap_mask(self, group: DataFrameGroupBy, next_from: pd.DataFrame):
        """
        Create a numpy mask to define row with gap

        Args:
            group: DataFrameGroupBy hole_id
            next_from: Dataframe next from values

        Returns: numpy mask for row with gap

        """
        mask = np.logical_and(
            next_from
            > group[self._cols.to_col],  # from is greater than current to : gap
            ~pd.isnull(
                next_from
            ),  # considers only shifted value (last shifted value is NaN, by definition)
        )
        return mask

    def _expansion_gap(
        self,
        df: pd.DataFrame,
        group: DataFrameGroupBy,
        hole_id: str,
        backward_expansion: bool,
    ) -> pd.DataFrame:
        """
        Insert row for gap (simple expansion)

        Args:
            df: input dataframe
            group: DataFrameGroupBy hole_id
            hole_id: group hole_id
            backward_expansion: True for backward expansion, False for forward expansion

        Returns: dataframe with new rows for gap

        """
        next_from = group[self._cols.from_col].shift(-1)
        mask = self._gap_mask(group, next_from)
        # New value is previous one for backward expansion, next one for forward expansion
        if backward_expansion:
            new_val = group[self._cols.val_cols]
        else:
            new_val = group[self._cols.val_cols].shift(-1)

        # Define added row
        newdf = pd.DataFrame(
            {
                self._cols.hole_id_col: hole_id,
                # The missing "from" is the previous "to"
                self._cols.from_col: group.loc[mask][self._cols.to_col],
                # The missing "to" is the next "from" (aka the shifted one)
                self._cols.to_col: next_from.loc[mask],
                # Update status
                self._cols.update_status_col: "expansion",
            }
            | {
                # values
                key: new_val[key].loc[mask]
                for key in self._cols.val_cols
            },
        )

        # Concat the new row(s) to the main DataFrame
        df = pd.concat([df, newdf])
        return df

    def _nearest_neighbor_gap(
        self, df: pd.DataFrame, group: DataFrameGroupBy, hole_id: str
    ) -> pd.DataFrame:
        """
        Insert rows for gap (nearest neighbor expansion)

        Args:
            df: input dataframe
            group: DataFrameGroupBy hole_id
            hole_id: group hole_id

        Returns: dataframe with new rows for gap

        """
        next_from = group[self._cols.from_col].shift(-1)
        next_val = group[self._cols.val_cols].shift(-1)

        mask = self._gap_mask(group, next_from)

        # Nearest neighbor : cur_to + (next_from - cur_to) / 2
        nearest_neighbor = (
            group.loc[mask][self._cols.to_col]
            + (next_from.loc[mask] - group.loc[mask][self._cols.to_col]) / 2.0
        )

        # Backward expansion
        backward_df = pd.DataFrame(
            {
                self._cols.hole_id_col: hole_id,
                # From : current to
                self._cols.from_col: group.loc[mask][self._cols.to_col],
                # To : nearest neighbor
                self._cols.to_col: nearest_neighbor,
                # Update status
                self._cols.update_status_col: "expansion",
            }
            | {
                # values
                key: group[key].loc[mask]
                for key in self._cols.val_cols
            },
        )

        # Forward expansion
        forward_df = pd.DataFrame(
            {
                self._cols.hole_id_col: hole_id,
                # From : nearest neighbor
                self._cols.from_col: nearest_neighbor,
                # To : next from
                self._cols.to_col: next_from.loc[mask],
                # Update status
                self._cols.update_status_col: "expansion",
            }
            | {
                # values
                key: next_val[key].loc[mask]
                for key in self._cols.val_cols
            },
        )

        # Concat the new rows to the main DataFrame
        df = pd.concat([df, backward_df, forward_df])
        return df

    def _rejection_gap(
        self, df: pd.DataFrame, group: DataFrameGroupBy, hole_id: str
    ) -> pd.DataFrame:
        """
        Remove all row with gap

        Args:
            df: input dataframe
            group: DataFrameGroupBy hole_id
            hole_id: group hole_id

        Returns: dataframe with remove of hole if gap

        """
        next_from = group[self._cols.from_col].shift(-1)
        mask = self._gap_mask(group, next_from)

        if np.count_nonzero(mask):
            df = df.drop(group.index)
        return df

    def _overlap_mask_with_prev_to(
        self, group: DataFrameGroupBy, prev_to: pd.DataFrame
    ):
        """
        Create a numpy mask to define row with overlap (get previous row)

        Args:
            group: DataFrameGroupBy hole_id
            prev_to: Dataframe previous to values

        Returns: numpy mask for row with overlap (get previous row)

        """
        mask = np.logical_and(
            prev_to
            > group[
                self._cols.from_col
            ],  # prev to is greater than current from : overlap
            ~pd.isnull(
                prev_to
            ),  # considers only shifted value (last shifted value is NaN, by definition)
        )
        return mask

    def _overlap_mask_with_next_from(
        self, group: DataFrameGroupBy, next_from: pd.DataFrame
    ):
        """
        Create a numpy mask to define row with overlap (get next row)

        Args:
            group: DataFrameGroupBy hole_id
            next_from: Dataframe next from values

        Returns: numpy mask for row with overlap (get next row)

        """
        mask = np.logical_and(
            next_from
            < group[
                self._cols.to_col
            ],  # current to is greater than next from : overlap
            ~pd.isnull(
                next_from
            ),  # considers only shifted value (last shifted value is NaN, by definition)
        )
        return mask

    def _expansion_overlap(
        self, df: pd.DataFrame, group: DataFrameGroupBy, backward_expansion: bool
    ) -> pd.DataFrame:
        """
        Update row for overlap (simple expansion)

        Args:
            df: input dataframe
            group:  DataFrameGroupBy hole_id
            backward_expansion: True for backward expansion, False for forward expansion

        Returns:  dataframe with updated rows for overlap

        """
        if backward_expansion:
            # Update value for row before overlap
            prev_to = group[self._cols.to_col].shift(1)
            mask = self._overlap_mask_with_prev_to(group, prev_to)
            # From = gap to
            df.loc[group.index, self._cols.from_col] = group[self._cols.from_col].mask(
                mask, prev_to
            )
            df.loc[group.index, self._cols.update_status_col] = group[
                self._cols.update_status_col
            ].mask(mask, "expansion")
        else:
            # Update value for row after overlap
            next_from = group[self._cols.from_col].shift(-1)
            mask = self._overlap_mask_with_next_from(group, next_from)
            # To = gap from
            df.loc[group.index, self._cols.to_col] = group[self._cols.to_col].mask(
                mask, next_from
            )
            df.loc[group.index, self._cols.update_status_col] = group[
                self._cols.update_status_col
            ].mask(mask, "expansion")

        return df

    def _nearest_neighbor_overlap(
        self, df: pd.DataFrame, group: DataFrameGroupBy
    ) -> pd.DataFrame:
        """
        Update rows for overlap (nearest neighbor expansion)

        Args:
            df: input dataframe
            group:  DataFrameGroupBy hole_id

        Returns:  dataframe with updated rows for overlap

        """
        prev_to = group[self._cols.to_col].shift(1)
        next_from = group[self._cols.from_col].shift(-1)

        current_to = group[self._cols.to_col]
        current_from = group[self._cols.from_col]

        mask_from = self._overlap_mask_with_prev_to(group, prev_to)
        mask_to = self._overlap_mask_with_next_from(group, next_from)

        # Update value for row before overlap
        df.loc[group.index, self._cols.to_col] = current_to.mask(
            mask_to, next_from + (current_to - next_from) / 2.0
        )

        # Update value for row after overlap
        df.loc[group.index, self._cols.from_col] = current_from.mask(
            mask_from, current_from + (prev_to - current_from) / 2.0
        )

        if np.count_nonzero(mask_to) and np.count_nonzero(mask_from):
            combine_mask = ma.mask_or(mask_to, mask_from)
            df.loc[group.index, self._cols.update_status_col] = group[
                self._cols.update_status_col
            ].mask(combine_mask, "expansion")

        return df

    def _rejection_overlap(
        self, df: pd.DataFrame, group: DataFrameGroupBy, hole_id: str
    ) -> pd.DataFrame:
        """
        Remove all row with overlap

        Args:
            df: input dataframe
            group: DataFrameGroupBy hole_id
            hole_id: group hole_id

        Returns: dataframe with remove of hole if overlap

        """
        prev_to = group[self._cols.to_col].shift(1)
        mask = self._overlap_mask_with_prev_to(group, prev_to)

        if np.count_nonzero(mask):
            df = df.drop(group.index)
        return df
