import pandas as pd
from tqdm import tqdm
from time import sleep
import warnings

from .csv_class import Csv
from .data_loader import Data_Loader

class CombinedDataset(Data_Loader):
    """
    A class for combining multiple datasets of a single type (Csv, Excel, etc.) into a single data loader

    Parameters
    ----------
    data_class: abc.ABCMeta 
        Data loader class of datasets to be combined
    loaders: list[Data_Loader]
        Individual data loader objects

    Methods
    -------
    CombinedDataset(data_class, url, datasets, *args, **kwargs)
        Constructor.
    load(*args, **kwargs)
        Load data for query. 
    get_count(*args, **kwargs)
        Get number of records/rows generated by query
    get_years(*args, **kwargs)
        Get years contained in data set
    """

    def __init__(self, data_class, url, datasets, *args, pbar=True, **kwargs):
        """CombinedDataset constructor

        Parameters
        ----------
        data_class : abc.ABCMeta 
            Class of data loader object
        url : str
            Base URL of dataset (will be combined with each value in datasets)
        datasets : str or List[str]
            Relative dataset URLs (each value will be combined with url). Can be a semi-colon separate string or a list.

        *args and **kwargs will be passed to the constructor of data_class
        """
        self.url = url
        self.data_class = data_class
        self.datasets = datasets
        
        self.loaders = []
        url = url[:-1] if url[-1]=='/' else url
        iter = tqdm(datasets, desc='Building Data Loaders', leave=False) if pbar else datasets
        for ds in iter:
            if isinstance(ds, list):
                # This case indicates tables that will be joined
                self.loaders.append(CombinedDataset(data_class, url, ds, *args, pbar=False, **kwargs))
                continue

            loc_kwargs = kwargs.copy()
            if 'url' in ds and 'raw.githubusercontent.com/openpolicedata/opd-datasets' in ds['url'] and ds['url'].endswith('.csv'):
                # This dataset has been re-posted on our GitHub page after being taken down by the original poster
                self.loaders.append(Csv(ds['url'], *args, **loc_kwargs))
            else:
                ds = ds.copy()
                cur_url = url + '/' + ds.pop('url') if 'url' in ds else url
                loc_kwargs['data_set'] = ds
                try:
                    self.loaders.append(data_class(cur_url, *args, **loc_kwargs))
                except ValueError as e:
                    if str(e)=='Excel file format cannot be determined, you must specify an engine manually.':
                        try:
                            # This may be a CSV file instead of an Excel file
                            self.loaders.append(Csv(cur_url, *args, **loc_kwargs))
                        except:
                            raise e
                        
                if ds!=datasets[-1]:
                    sleep(0.5)  # Reduce likelihood of timeout due to repeated requests



    def isfile(self):
        '''Returns True to indicate that Csv data is file-based

        Returns
        -------
        True
        '''
        return True
    
    def load(self, nrows=None, offset=None, pbar=True, **kwargs):
        """Load data for query. 

        **kwargs will be passed to the load function of the data_class
        """

        dfs = []
        on = []
        date_warn = force_subject_warn = force_officer_warn = False
        first_time = True
        if '_first_time' in kwargs:
            kwargs.pop('_first_time')
        iter = tqdm(self.loaders, desc='Loading data files', leave=False) if pbar else self.loaders
        for k, loader in enumerate(iter):
            if isinstance(self.datasets[k],list):
                # Tables in dfs will be merged
                on.append(self.datasets[k][0]['on'])

            dfs.append(loader.load(_first_time=first_time, **kwargs))
            first_time = False
            if loader!=self.loaders[-1]:
                sleep(0.5)  # Reduce likelihood of timeout due to repeated requests

            if 'www.albemarle.org' in loader.url:
                # Renamed
                if 'Stop Date' in dfs[-1] and 'Date' in dfs[0]: # Column renamed
                    dfs[-1] = dfs[-1].rename(columns={'Stop Date':'Date'})
                    if not date_warn:
                        date_warn = True
                        warnings.warn("Renaming date column because name of column names changes in some of the monthly data files")
                if 'Force Used by Subject' in dfs[-1] and 'Physical Force by Subject' in dfs[0]: # Column renamed
                    dfs[-1] = dfs[-1].rename(columns={'Force Used by Subject':'Physical Force by Subject'})
                    if not force_subject_warn:
                        force_subject_warn = True
                        warnings.warn("Renaming force by subject column because name of column names changes in some of the monthly data files")
                if 'Force Used by Officer' in dfs[-1] and 'Physical Force by Officer' in dfs[0]: # Column renamed
                    dfs[-1] = dfs[-1].rename(columns={'Force Used by Officer':'Physical Force by Officer'})
                    if not force_officer_warn:
                        force_officer_warn = True
                        warnings.warn("Renaming force by officer column because name of column names changes in some of the monthly data files")

        if len(on)==0:
            df = pd.concat(dfs, ignore_index=True)
        else:
            raise NotImplementedError("This code currently does not affect any dataset and the result should be reviewed prior to usage")
            df = dfs[0]
            for k in range(1,len(dfs)):
                df = df.merge(dfs[k], how='outer', left_on=on[0], right_on=on[k])

        if offset!=None:
            df = df.iloc[offset:]
        if nrows!=None:
            df = df.head(nrows)

        return df
        
    def get_count(self, *args, **kwargs):   
        """Get number of records/rows generated by query

        *args and **kwargs will be passed to the load function of the data_class
        """

        count = 0
        first_time = True
        for loader in self.loaders:
            count+=loader.get_count(*args, _first_time=first_time, **kwargs)
            first_time = False

        return count
    
    
    def get_years(self, *args, **kwargs):   
        """Get years contained in data set

        *args and **kwargs will be passed to the load function of the data_class
        """

        years = []
        for loader in self.loaders:
            years.extend(loader.get_years(*args, **kwargs))

        return years

