"""
Copyright 2022 Ilia Moiseev

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from typing import List, Dict
import pandas as pd
from dask import dataframe as dd

from ..meta import AggregateValidator, DataValidationException
from ..data import Dataset, Modifier, Iterator, SequentialCacher


class TableDataset(Dataset):
    """
    Wrapper for `pd.DataFrame`s
    """
    def __init__(self, *args, t=None, **kwargs):
        """
        Parameters
        ----------
        t:
            pd.DataFrame or TableDataset to be set as table
        """
        super().__init__(*args, **kwargs)
        if isinstance(t, pd.DataFrame):
            self._table = t
        elif isinstance(t, TableDataset):
            self._table = t._table
        elif t is None:
            self._table = pd.DataFrame()
        else:
            raise TypeError('Input table is not a pandas.DataFrame nor TableDataset')

    def __getitem__(self, index):
        """
        Returns row from table by index
        """
        return self._table.iloc[index]

    def __repr__(self):
        return f'{super().__repr__()}\n {repr(self._table)}'

    def __len__(self):
        """
        Return len of the table
        """
        return len(self._table)

    def get_meta(self) -> List[Dict]:
        meta = super().get_meta()
        meta[0].update({
            'name': repr(self),
            'columns': list(self._table.columns),
            'len': len(self),
            'info': self._table.describe().to_dict()
        })
        return meta

    def to_csv(self, path, **kwargs):
        """
        Saves the table to .csv
        """
        self._table.to_csv(path, **kwargs)


class TableFilter(TableDataset, Modifier):
    """
    Filter for table values
    """
    def __init__(self, dataset, mask, *args, **kwargs):
        """
        Parameters
        ----------
        dataset: TableDataset
            Dataset to be filtered
        mask: Iterable[bool]
            Binary mask to select values from table
        """
        super().__init__(dataset, t=dataset._table, *args, **kwargs)
        init_len = len(dataset)

        self._table = self._table[mask]
        print(f'Length before filtering: {init_len}, length after: {len(self._table)}')


class CSVDataset(TableDataset):
    """
    Wrapper for .csv files.
    """
    def __init__(self, csv_file_path, *args, **kwargs):
        """
        Passes all args and kwargs to the read_csv

        Parameters
        ----------
        csv_file_path:
            path to the .csv file
        """
        t = pd.read_csv(csv_file_path, *args, **kwargs)
        super().__init__(t=t, **kwargs)


class PartedTableLoader(Dataset):
    """
    Works like CSVDataset, but uses dask to load tables
    and returns partitions on __getitem__
    """
    def __init__(self, csv_file_path, *args, **kwargs):
        super().__init__(**kwargs)
        self._table = dd.read_csv(csv_file_path, *args, **kwargs)

    def __getitem__(self, index):
        """
        Returns partition under the index
        """
        return self._table.get_partition(index).compute()

    def __len__(self):
        """
        The number of partitions
        """
        return self._table.npartitions


class TableIterator(Iterator):
    """
    Iterates over the table from path by the chunks.
    """
    def __init__(self, csv_file_path, *args, chunk_size=1000, **kwargs):
        """
        Parameters
        ----------
        csv_file_path:
            path to the .csv file

        chunk_size: int
            number of rows to return in one __next__
        """
        self.chunk_size = chunk_size
        super().__init__(pd.read_csv(csv_file_path,
                         iterator=True, *args, **kwargs))

    def __next__(self):
        return self._data.get_chunk(self.chunk_size)


class LargeCSVDataset(SequentialCacher):
    """
    SequentialCacher over large .csv file.
    Loads table by partitions.
    """
    def __init__(self, csv_file_path, *args, **kwargs):
        dataset = PartedTableLoader(csv_file_path, *args, **kwargs)
        self.num_batches = dataset._table.npartitions
        self.bs = self.len // self.num_batches

        super().__init__(dataset, self.bs)
        self.len = len(dataset._table)

    def _load(self, index):
        del self.batch
        self.batch = TableDataset(self._dataset[index])

    def __len__(self):
        return self.len


class NullValidator(TableDataset, AggregateValidator):
    """
    Checks there are no null values in the table.
    """
    def __init__(self, dataset: TableDataset, *args, **kwargs) -> None:
        super().__init__(dataset, self.check_nulls,
                         *args, t=dataset._table, **kwargs)

    def check_nulls(self, x):
        mask = x._table.isnull().values
        if ~(mask.any()):
            return True
        else:
            total = mask.sum()
            by_columns = mask.sum(axis=0)
            missing = pd.DataFrame(by_columns.reshape(1, len(by_columns)), columns=x._table.columns)
            raise DataValidationException(
                f'There were NaN-values in {repr(self._dataset)}\n'
                f'Total count: {total}\n'
                f'By columns:\n'
                f'{missing}'
            )
