import numpy as np
from abc import ABC, abstractmethod
import pprint
from typing import Tuple

from ada.data_containers._base import _ActiData, _Epoched, _Raw
from ada.data_containers.generic import GenericData
from ada.data_containers.scored import ScoredShort


np.set_printoptions(legacy='1.21')


class _Long(ABC):

    @abstractmethod
    def fit(self, data: _ActiData, ch_name: str | None = None):
        pass

    def _get_data(self, data: _ActiData, ch_name: str | None = None) -> np.ndarray:
        if ch_name is None:
            if isinstance(data, (_Raw, _Epoched, GenericData)):
                return data.to_score
            elif isinstance(data, ScoredShort):
                return data.score
            else:
                raise ValueError("Invalid data type")
        else:
            return data.data[data.channel_names.index(ch_name)]


class _LongContainer(ABC):

    _fit_params: dict = {}

    def __repr__(self):
        return pprint.pformat(self._fit_params, indent=4)
    
    @staticmethod
    @abstractmethod
    def load_file(path: str) -> "_LongContainer":
        pass

    @staticmethod
    @abstractmethod
    def _construct_from_load(data: np.ndarray, fit: dict) -> "_LongContainer":
        pass
    
    @abstractmethod
    def plot(self, out_path: str | None = None):
        pass

    @abstractmethod
    def export(self, out_path: str):
        pass

    @property
    @abstractmethod
    def id(self) -> str:
        pass

    @staticmethod
    def _generic_export(path: str, to_export: list):
        print('Saving data to file: {}'.format(path))
        if path[-8:] != 'ada.long':
            if path[-3:] == 'ada':
                path = path + '.long'
            else:
                path = path + '.ada.long'
        with open(path, 'wb') as f:
            np.savez_compressed(f, *to_export)
        print('Done saving to {}'.format(path))

    @staticmethod
    def _generic_load(path: str | list) -> Tuple[np.ndarray, dict]:
        def arr_no(n):
            return "arr_{}".format(n)
        
        if isinstance(path, str):
            files = np.load(path, allow_pickle=True)
            data = files[arr_no(0)]
            fit_params = files[arr_no(1)].item()
            return data, fit_params
        return path[0], path[1]
        

    @abstractmethod
    def save_csv(self, out_path: str):
        pass

    @property
    def fit_params(self) -> dict:
        """Parameters of the fitted model."""
        return self._fit_params
