# -*- coding: utf-8 -*-
"""
General Joint Probabilty Distribution calculator
"""
from concurrent.futures import as_completed
import gc
import logging
import h5py
import numpy as np
import os
import pandas as pd
from warnings import warn

from rex.version import __version__
from rex.renewable_resource import WindResource
from rex.resource import Resource
from rex.utilities.execution import SpawnProcessPool
from rex.utilities.loggers import log_mem, log_versions
from rex.utilities.utilities import slice_sites, to_records_array

logger = logging.getLogger(__name__)


class JointPD:
    """
    Compute the joint probability distribution between the desired variables
    """
    def __init__(self, res_h5, res_cls=Resource, hsds=False):
        """
        Parameters
        ----------
        res_h5 : str
            Path to resource h5 file(s)
        res_cls : Class, optional
            Resource handler class to use to access res_h5,
            by default Resource
        hsds : bool, optional
            Boolean flag to use h5pyd to handle .h5 'files' hosted on AWS
            behind HSDS, by default False
        """
        log_versions(logger)
        self._res_h5 = res_h5
        self._res_cls = res_cls
        self._hsds = hsds

    @property
    def res_h5(self):
        """
        Path to resource h5 file(s)

        Returns
        -------
        str
        """
        return self._res_h5

    @property
    def res_cls(self):
        """
        Resource class to use to access wind_h5

        Returns
        -------
        Class
        """
        return self._res_cls

    @staticmethod
    def _make_bins(start, stop, step):
        """
        Create bin edges from bin range

        Parameters
        ----------
        start : int
            bin range start value
        stop : int
            bin range stop value
        step : int
            bin range step value

        Returns
        -------
        bin_edges : ndarray
            Vector of inclusive bin edges
        """
        bin_edges = np.arange(start, stop + step, step)

        return bin_edges

    @classmethod
    def compute_joint_pd(cls, res_h5, dset1, dset2, bins1, bins2,
                         res_cls=Resource, hsds=False,
                         sites_slice=None):
        """
        Compute the joint probability distribution between the two given
        datasets using the given bins for given sites

        Parameters
        ----------
        res_h5 : str
            Path to resource h5 file(s)
        dset1 : str
            Dataset 1 to generate joint probability distribution for
        dset2 : str
            Dataset 2 to generate joint probabilty distribution for
        bins1 : tuple
            (start, stop, step) for dataset 1 bins
        bins2 : tuple
            (start, stop, step) for dataset 2 bins
        sites : list | slice, optional
        res_cls : Class, optional
            Resource handler class to use to access res_h5,
            by default Resource
        hsds : bool, optional
            Boolean flag to use h5pyd to handle .h5 'files' hosted on AWS
            behind HSDS, by default False
        sites_slice : slice | None, optional
            Sites to extract, if None all, by default None

        Returns
        -------
        jpd : dict
            Dictionary of joint probabilty distribution densities for given
            sites
        """
        if sites_slice is None:
            sites_slice = slice(None, None, None)
        elif isinstance(sites_slice, int):
            sites_slice = [sites_slice]

        with res_cls(res_h5, hsds=hsds) as f:
            arr1 = f[dset1, :, sites_slice]
            arr2 = f[dset2, :, sites_slice]

        bins1 = cls._make_bins(*bins1)
        bins2 = cls._make_bins(*bins2)

        if isinstance(sites_slice, slice) and sites_slice.stop:
            gids = list(range(*sites_slice.indices(sites_slice.stop)))
        elif isinstance(sites_slice, (list, np.ndarray)):
            gids = sites_slice

        jpd = {}
        for i, (a1, a2) in enumerate(zip(arr1.T, arr2.T)):
            jpd[gids[i]] = np.histogram2d(a1, a2,
                                          bins=(bins1, bins2),
                                          density=True)[0].astype(np.float32)

        return jpd

    def _get_slices(self, dset1, dset2, sites=None, chunks_per_slice=5):
        """
        Get slices to extract, ensure the shapes of dset1 and 2 match.

        Parameters
        ----------
        dset1 : str
            Dataset 1 to generate joint probability distribution for
        dset2 : str
            Dataset 2 to generate joint probabilty distribution for
        sites : list | slice, optional
            Subset of sites to extract, by default None or all sites
        chunks_per_slice : int, optional
            Number of chunks to extract in each slice, by default 5

        Returns
        -------
        slices : list
            List of slices to extract
        """
        with self.res_cls(self.res_h5) as f:
            shape, _, chunks = f.get_dset_properties(dset1)
            shape2, _, _ = f.get_dset_properties(dset2)

        if shape != shape2:
            msg = ("The shape of {}: {}, does not match the shape of {}: {}!"
                   .format(dset1, shape, dset2, shape2))
            logger.error(msg)
            raise RuntimeError(msg)

        slices = slice_sites(shape, chunks, sites=sites,
                             chunks_per_slice=chunks_per_slice)

        return slices

    def compute(self, dset1, dset2, bins1, bins2, sites=None, max_workers=None,
                chunks_per_worker=5):
        """
        Compute joint probability distribution between given datasets using
        given bins for all sites.

        Parameters
        ----------
        dset1 : str
            Dataset 1 to generate joint probability distribution for
        dset2 : str
            Dataset 2 to generate joint probabilty distribution for
        bins1 : tuple
            (start, stop, step) for dataset 1 bins
        bins2 : tuple
            (start, stop, step) for dataset 2 bins
        sites : list | slice, optional
            Subset of sites to extract, by default None or all sites
        max_workers : None | int, optional
            Number of workers to use, if 1 run in serial, if None use all
            available cores, by default None
        chunks_per_worker : int, optional
            Number of chunks to extract on each worker, by default 5

        Returns
        -------
        jpd: pandas.DataFrame
            DataFrame of joint probability distribution between given datasets
            with given bins
        """
        if max_workers is None:
            max_workers = os.cpu_count()

        slices = self._get_slices(dset1, dset2, sites,
                                  chunks_per_slice=chunks_per_worker)
        if len(slices) == 1:
            max_workers = 1

        jpd = {}
        if max_workers > 1:
            msg = ('Computing the joint probability distribution between {} '
                   'and {} in parallel using {} workers'
                   .format(dset1, dset2, max_workers))
            logger.info(msg)

            loggers = [__name__, 'rex']
            with SpawnProcessPool(max_workers=max_workers,
                                  loggers=loggers) as exe:
                futures = []
                for sites_slice in slices:
                    future = exe.submit(self.compute_joint_pd,
                                        self.res_h5, dset1, dset2,
                                        bins1, bins2,
                                        res_cls=self.res_cls,
                                        hsds=self._hsds,
                                        sites_slice=sites_slice)
                    futures.append(future)

                for i, future in enumerate(as_completed(futures)):
                    jpd.update(future.result())
                    logger.debug('Completed {} out of {} workers'
                                 .format((i + 1), len(futures)))

        else:
            msg = ('Computing the joint probability distribution between {} '
                   'and {} in serial.'
                   .format(dset1, dset2))
            logger.info(msg)
            for i, sites_slice in enumerate(slices):
                jpd.update(self.compute_joint_pd(
                    self.res_h5, dset1, dset2,
                    bins1, bins2,
                    res_cls=self.res_cls,
                    hsds=self._hsds,
                    sites_slice=sites_slice))
                logger.debug('Completed {} out of {} sets of sites'
                             .format((i + 1), len(slices)))

        gc.collect()
        log_mem(logger)
        bins1 = self._make_bins(*bins1)
        bins2 = self._make_bins(*bins2)
        index = np.meshgrid(bins1[:-1], bins2[:-1], indexing='ij')
        index = np.array(index).T.reshape(-1, 2).astype(np.int16)
        index = pd.MultiIndex.from_arrays(index.T, names=(dset1, dset2))
        jpd = pd.DataFrame({k: v.flatten(order='F') for k, v
                            in jpd.items()}, index=index).sort_index(axis=1)

        return jpd

    def save(self, jpd, out_fpath):
        """
        Save joint probability distribution to disk

        Parameters
        ----------
        jpd : pandas.DataFrame
            Table of joint probability distribution densities to save
        out_fpath : str
            .csv, or .h5 file path to save joint probability
            distribution to
        """
        with self.res_cls(self.res_h5) as f:
            meta = f['meta', jpd.columns.values]

        logger.info('Writing joint probability distribution to {}'
                    .format(out_fpath))
        if out_fpath.endswith('.csv'):
            jpd.to_csv(out_fpath)
            meta_fpath = out_fpath.split('.')[0] + '_meta.csv'
            if os.path.exists(meta_fpath):
                msg = ("Site meta data already exists at {}!")
                logger.warning(msg)
                warn(msg)
            else:
                logger.debug('Writing site meta data to {}'
                             .format(meta_fpath))
                meta.to_csv(meta_fpath, index=False)
        elif out_fpath.endswith('.h5'):
            with h5py.File(out_fpath, mode='w') as f:
                f.attrs['rex version'] = __version__
                for i, n in enumerate(jpd.index.names):
                    logger.info('')
                    data = np.array(jpd.index.get_level_values(i))
                    dset = '{}-bins'.format(n)
                    logger.debug('Writing {}'.format(dset))
                    f.create_dataset(dset, shape=data.shape, dtype=data.dtype,
                                     data=data)

                logger.debug('Writing joint probability density values to jpd')
                data = jpd.values
                f.create_dataset('jpd', shape=data.shape, dtype=data.dtype,
                                 data=data)

                logger.debug('Writing site meta data to meta')
                meta = to_records_array(meta)
                f.create_dataset('meta', shape=meta.shape, dtype=meta.dtype,
                                 data=meta)
        else:
            msg = ("Cannot save joint probability distribution, expecting "
                   ".csv or .h5 path, but got: {}".format(out_fpath))
            logger.error(msg)
            raise OSError(msg)

    @staticmethod
    def plot_joint_pd(jpd, site=None, **kwargs):
        """
        Plot the mean joint probability distribution accross all sites
        (site=None), or the distribution for the single given site

        Parameters
        ----------
        jpd: pandas.DataFrame
            DataFrame of joint probability distribution between given datasets
            with given bins
        site : int, optional
            Site to plot distribution for, if None plot mean distribution
            across all sites, by default None
        """
        x, y = jpd.index.names
        if site is not None:
            msg = ("Can only plot the joint probabilty distribution for a "
                   "single site or the mean probability distribution accross "
                   "all sites (site=None), you provided: {}".format(site))
            assert isinstance(site), msg
            plt = jpd.loc[:, [site]].reset_index()
        else:
            site = 'mean_jpd'
            plt = jpd.mean(axis=1)
            plt.name = site
            plt = plt.reset_index()

        plt.plot.scatter(x=x, y=y, c=site, **kwargs)

    @classmethod
    def run(cls, res_h5, dset1, dset2, bins1, bins2,
            sites=None, res_cls=Resource, hsds=False,
            max_workers=None, chunks_per_worker=5, out_fpath=None):
        """
        Compute joint probability distribution between given datasets using
        given bins

        Parameters
        ----------
        res_h5 : str
            Path to resource h5 file(s)
        dset1 : str
            Dataset 1 to generate joint probability distribution for
        dset2 : str
            Dataset 2 to generate joint probabilty distribution for
        bins1 : tuple
            (start, stop, step) for dataset 1 bins
        bins2 : tuple
            (start, stop, step) for dataset 2 bins
        sites : list | slice, optional
            Subset of sites to extract, by default None or all sites
        res_cls : Class, optional
            Resource class to use to access res_h5, by default Resource
        hsds : bool, optional
            Boolean flag to use h5pyd to handle .h5 'files' hosted on AWS
            behind HSDS, by default False
        max_workers : None | int, optional
            Number of workers to use, if 1 run in serial, if None use all
            available cores, by default None
        chunks_per_worker : int, optional
            Number of chunks to extract on each worker, by default 5
        out_fpath : str, optional
            .csv, or .h5 file path to save joint probability
            distribution to

        Returns
        -------
        out : pandas.DataFrame
            DataFrame of joint probability distribution between given datasets
            with given bins
        """
        logger.info('Computing joint probability distribution between {} and '
                    '{} in {}'
                    .format(dset1, dset2, res_h5))
        logger.debug('Computing joint probability distribution using:'
                     '\n-{} bins: {}'
                     '\n-{} bins: {}'
                     '\n-max workers: {}'
                     '\n-chunks per worker: {}'
                     .format(dset1, bins1, dset2, bins2, max_workers,
                             chunks_per_worker))
        jpd = cls(res_h5, res_cls=res_cls, hsds=hsds)
        out = jpd.compute(dset1, dset2, bins1, bins2,
                          sites=sites,
                          max_workers=max_workers,
                          chunks_per_worker=chunks_per_worker)
        if out_fpath is not None:
            jpd.save(out, out_fpath)

        return out

    @classmethod
    def wind_rose(cls, wind_h5, hub_height, wspd_bins=(0, 30, 1),
                  wdir_bins=(0, 360, 5), sites=None, res_cls=WindResource,
                  hsds=False, max_workers=None, chunks_per_worker=5,
                  out_fpath=None):
        """
        Compute wind rose at given hub height

        Parameters
        ----------
        wind_h5 : str
            Path to resource h5 file(s)
        hub_height : str | int
            Hub-height to compute wind rose at
        wspd_bins : tuple
            (start, stop, step) for wind speed bins
        wdir_bins : tuple
            (start, stop, step) for wind direction bins
        sites : list | slice, optional
            Subset of sites to extract, by default None or all sites
        res_cls : Class, optional
            Resource class to use to access wind_h5, by default Resource
        hsds : bool, optional
            Boolean flag to use h5pyd to handle .h5 'files' hosted on AWS
            behind HSDS, by default False
        max_workers : None | int, optional
            Number of workers to use, if 1 run in serial, if None use all
            available cores, by default None
        chunks_per_worker : int, optional
            Number of chunks to extract on each worker, by default 5
        out_fpath : str, optional
            .csv, or .h5 file path to save wind rose to

        Returns
        -------
        wind_rose : pandas.DataFrame
            DataFrame of wind rose frequencies at desired hub-height
        """
        logger.info('Computing wind rose for {}m wind in {}'
                    .format(hub_height, wind_h5))
        logger.debug('Computing wind rose using:'
                     '\n-wind speed bins: {}'
                     '\n-wind direction bins: {}'
                     '\n-max workers: {}'
                     '\n-chunks per worker: {}'
                     .format(wspd_bins, wdir_bins, max_workers,
                             chunks_per_worker))
        wind_rose = cls(wind_h5, res_cls=res_cls, hsds=hsds)
        wspd_dset = 'windspeed_{}m'.format(hub_height)
        wdir_dset = 'winddirection_{}m'.format(hub_height)
        out = wind_rose.compute(wspd_dset, wdir_dset, wspd_bins, wdir_bins,
                                sites=sites,
                                max_workers=max_workers,
                                chunks_per_worker=chunks_per_worker)
        if out_fpath is not None:
            wind_rose.save(out, out_fpath)

        return out
