#  Copyright (c) 2023 Mira Geoscience Ltd.
#
#  This file is part of geoh5py.
#
#  geoh5py is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  geoh5py is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public License
#  along with geoh5py.  If not, see <https://www.gnu.org/licenses/>.
from __future__ import annotations

import uuid
from abc import ABC, abstractmethod
from typing import cast

import numpy as np

from ...data import Data, ReferencedData
from ..curve import Curve
from ..object_type import ObjectType


class BaseElectrode(Curve, ABC):
    _potential_electrodes: PotentialElectrode | None = None
    _current_electrodes: CurrentElectrode | None = None

    def __init__(self, object_type: ObjectType, **kwargs):
        self._metadata: dict | None = None
        self._ab_cell_id: ReferencedData | None = None

        super().__init__(object_type, **kwargs)

    @property
    def ab_cell_id(self) -> ReferencedData | None:
        """
        Reference data entity mapping cells to a unique current dipole.
        """
        if getattr(self, "_ab_cell_id", None) is None:
            child = self.get_data("A-B Cell ID")
            if any(child) and isinstance(child[0], ReferencedData):
                self.ab_cell_id = child[0]

        if getattr(self, "_ab_cell_id", None) is not None:
            return self._ab_cell_id

        return None

    @ab_cell_id.setter
    def ab_cell_id(self, data: Data | np.ndarray):
        if isinstance(data, Data):
            if not isinstance(data, ReferencedData):
                raise TypeError(f"ab_cell_id must be of type {ReferencedData}")

            if data.parent.uid == self.uid:
                self._ab_cell_id = data
            else:
                self._ab_cell_id = cast(ReferencedData, data.copy(parent=self))
        else:
            if data.dtype != np.int32:
                print("ab_cell_id values will be converted to type 'int32'")

            if any(self.get_data("A-B Cell ID")):
                child = self.get_data("A-B Cell ID")[0]
                if isinstance(child, ReferencedData):
                    child.values = data.astype(np.int32)
            else:
                complement: CurrentElectrode | PotentialElectrode = (
                    self.current_electrodes
                    if isinstance(self, PotentialElectrode)
                    else self.potential_electrodes
                )

                if complement is not None and complement.ab_cell_id is not None:
                    entity_type = complement.ab_cell_id.entity_type
                else:
                    value_map = {ii: str(ii) for ii in range(data.max() + 1)}
                    value_map[0] = "Unknown"
                    entity_type = {  # type: ignore
                        "primitive_type": "REFERENCED",
                        "value_map": value_map,
                    }

                data = self.add_data(
                    {
                        "A-B Cell ID": {
                            "values": data.astype(np.int32),
                            "association": "CELL",
                            "entity_type": entity_type,
                        }
                    }
                )

                if isinstance(data, ReferencedData):
                    self._ab_cell_id = data

    @property
    def ab_map(self) -> dict | None:
        """
        Get the ReferenceData.value_map of the ab_value_id
        """
        if isinstance(self.ab_cell_id, ReferencedData):
            return self.ab_cell_id.value_map
        return None

    def copy(
        self,
        parent=None,
        copy_children: bool = True,
        clear_cache: bool = False,
        mask: np.ndarray | None = None,
        cell_mask: np.ndarray | None = None,
        **kwargs,
    ):
        """
        Sub-class extension of :func:`~geoh5py.objects.cell_object.CellObject.copy`.
        """
        if parent is None:
            parent = self.parent

        omit_list = [
            "_ab_cell_id",
            "_metadata",
            "_potential_electrodes",
            "_current_electrodes",
        ]
        new_entity = super().copy(
            parent=parent,
            clear_cache=clear_cache,
            copy_children=copy_children,
            mask=mask,
            cell_mask=cell_mask,
            omit_list=omit_list,
            **kwargs,
        )

        if self.cells is not None:
            if mask is not None:
                cell_mask = np.all(mask[self.cells], axis=1)
            else:
                cell_mask = np.ones(self.cells.shape[0], dtype=bool)

            if self.ab_cell_id is not None and self.ab_cell_id.values is not None:
                new_entity.ab_cell_id = self.ab_cell_id.values[cell_mask]

        complement: CurrentElectrode | PotentialElectrode = (
            self.current_electrodes
            if isinstance(self, PotentialElectrode)
            else self.potential_electrodes
        )

        # Set the mask of the complement
        if (
            new_entity.ab_cell_id is not None
            and complement is not None
            and complement.ab_cell_id is not None
            and complement.ab_cell_id.values is not None
            and complement.vertices is not None
            and complement.cells is not None
        ):
            intersect = np.intersect1d(
                new_entity.ab_cell_id.values,
                complement.ab_cell_id.values,
            )
            cell_mask = np.r_[
                [(val in intersect) for val in complement.ab_cell_id.values]
            ]

            # Convert cell indices to vertex indices
            mask = np.zeros(complement.vertices.shape[0], dtype=bool)
            mask[complement.cells[cell_mask, :]] = True

            new_complement = super(Curve, complement).copy(  # type: ignore
                parent=parent,
                omit_list=omit_list,
                copy_children=copy_children,
                clear_cache=clear_cache,
                mask=mask,
                cell_mask=cell_mask,
            )

            if isinstance(self, PotentialElectrode):
                new_entity.current_electrodes = new_complement
            else:
                new_entity.potential_electrodes = new_complement

            if new_complement.ab_cell_id is None and complement.ab_cell_id is not None:
                new_complement.ab_cell_id = complement.ab_cell_id.values[cell_mask]

            # Re-number the ab_cell_id
            value_map = {
                val: ind
                for ind, val in enumerate(
                    np.r_[0, np.unique(new_entity.current_electrodes.ab_cell_id.values)]
                )
            }
            new_map = {
                val: new_entity.current_electrodes.ab_cell_id.value_map.map[val]
                for val in value_map.values()
            }
            new_complement.ab_cell_id.values = np.asarray(
                [value_map[val] for val in new_complement.ab_cell_id.values]
            )
            new_entity.ab_cell_id.values = np.asarray(
                [value_map[val] for val in new_entity.ab_cell_id.values]
            )
            new_entity.ab_cell_id.value_map.map = new_map

        return new_entity

    @property
    @abstractmethod
    def current_electrodes(self):
        """
        The associated current_electrodes (transmitters)
        """

    @property
    def metadata(self):
        """
        Metadata attached to the entity.
        """
        if getattr(self, "_metadata", None) is None:
            metadata = self.workspace.fetch_metadata(self.uid)
            self._metadata = metadata
        return self._metadata

    @metadata.setter
    def metadata(self, values):
        if not len(values) == 2:
            raise ValueError(
                f"Metadata must have two key-value pairs. {values} provided."
            )

        default_keys = ["Current Electrodes", "Potential Electrodes"]

        if list(values.keys()) != default_keys:
            raise ValueError(f"Input metadata must have for keys {default_keys}")

        if self.workspace.get_entity(values["Current Electrodes"])[0] is None:
            raise IndexError("Input Current Electrodes uuid not present in Workspace")

        if self.workspace.get_entity(values["Potential Electrodes"])[0] is None:
            raise IndexError("Input Potential Electrodes uuid not present in Workspace")

        self._metadata = values
        self.workspace.update_attribute(self, "metadata")

    @property
    @abstractmethod
    def potential_electrodes(self):
        """
        The associated potential_electrodes (receivers)
        """


class PotentialElectrode(BaseElectrode):
    """
    Ground potential electrode (receiver).
    """

    __TYPE_UID = uuid.UUID("{275ecee9-9c24-4378-bf94-65f3c5fbe163}")

    @property
    def current_electrodes(self):
        """
        The associated current electrode object (sources).
        """
        if getattr(self, "_current_electrodes", None) is None:
            if self.metadata is not None and "Current Electrodes" in self.metadata:
                transmitter = self.metadata["Current Electrodes"]
                transmitter_entity = self.workspace.get_entity(transmitter)[0]

                if isinstance(transmitter_entity, CurrentElectrode):
                    self._current_electrodes = transmitter_entity

        return self._current_electrodes

    @current_electrodes.setter
    def current_electrodes(self, current_electrodes: CurrentElectrode):
        if not isinstance(current_electrodes, CurrentElectrode):
            raise TypeError(
                f"Provided current_electrodes must be of type {CurrentElectrode}. "
                f"{type(current_electrodes)} provided."
            )

        metadata = {
            "Current Electrodes": current_electrodes.uid,
            "Potential Electrodes": self.uid,
        }

        self.metadata = metadata
        current_electrodes.metadata = metadata

        if isinstance(current_electrodes.ab_cell_id, ReferencedData) and isinstance(
            self.ab_cell_id, ReferencedData
        ):
            self.ab_cell_id.entity_type = current_electrodes.ab_cell_id.entity_type

    @property
    def potential_electrodes(self):
        """
        The associated potential_electrodes (receivers)
        """
        return self

    @classmethod
    def default_type_uid(cls) -> uuid.UUID:
        """
        :return: Default unique identifier
        """
        return cls.__TYPE_UID


class CurrentElectrode(BaseElectrode):
    """
    Ground direct current electrode (transmitter).
    """

    __TYPE_UID = uuid.UUID("{9b08bb5a-300c-48fe-9007-d206f971ea92}")

    def __init__(self, object_type: ObjectType, **kwargs):
        self._current_line_id: uuid.UUID | None = None

        super().__init__(object_type, **kwargs)

    @classmethod
    def default_type_uid(cls) -> uuid.UUID:
        """
        :return: Default unique identifier
        """
        return cls.__TYPE_UID

    @property
    def current_electrodes(self):
        """
        The associated current electrode object (sources).
        """
        return self

    @current_electrodes.setter
    def current_electrodes(self, _):
        ...

    @property
    def potential_electrodes(self) -> PotentialElectrode | None:
        """
        The associated potential_electrodes (receivers)
        """
        if getattr(self, "_potential_electrodes", None) is None:
            if self.metadata is not None and "Potential Electrodes" in self.metadata:
                potential = self.metadata["Potential Electrodes"]
                potential_entity = self.workspace.get_entity(potential)[0]

                if isinstance(potential_entity, PotentialElectrode):
                    self._potential_electrodes = potential_entity

        return self._potential_electrodes

    @potential_electrodes.setter
    def potential_electrodes(self, potential_electrodes: PotentialElectrode):
        if not isinstance(potential_electrodes, PotentialElectrode):
            raise TypeError(
                f"Provided potential_electrodes must be of type {PotentialElectrode}. "
                f"{type(potential_electrodes)} provided."
            )

        metadata = {
            "Current Electrodes": self.uid,
            "Potential Electrodes": potential_electrodes.uid,
        }

        self.metadata = metadata
        potential_electrodes.metadata = metadata

        if isinstance(potential_electrodes.ab_cell_id, ReferencedData) and isinstance(
            self.ab_cell_id, ReferencedData
        ):
            potential_electrodes.ab_cell_id.entity_type = self.ab_cell_id.entity_type

    def add_default_ab_cell_id(self):
        """
        Utility function to set ab_cell_id's based on curve cells.
        """
        if getattr(self, "cells", None) is None or self.n_cells is None:
            raise AttributeError(
                "Cells must be set before assigning default ab_cell_id"
            )

        data = np.arange(self.n_cells) + 1
        value_map = {ii: str(ii) for ii in range(self.n_cells + 1)}
        value_map[0] = "Unknown"
        ab_cell_id = self.add_data(
            {
                "A-B Cell ID": {
                    "values": data,
                    "association": "CELL",
                    "entity_type": {
                        "primitive_type": "REFERENCED",
                        "value_map": value_map,
                    },
                }
            }
        )
        if isinstance(ab_cell_id, ReferencedData):
            ab_cell_id.entity_type.name = "A-B"
            self._ab_cell_id = ab_cell_id
