"""
Experimental module exposing variants of the HiC3DeFDR model for benchmarking
purposes.
"""

import numpy as np
import scipy.stats as stats

from hic3defdr.analysis import HiC3DeFDR
from hic3defdr.util.logging import eprint
from hic3defdr.util.dispersion import mme_per_pixel
from hic3defdr.util.clusters import load_clusters
from hic3defdr.util.parallelization import parallel_apply
import hic3defdr.util.dispersion as dispersion


def poisson_fit_mu_hat(raw, f):
    return np.average(raw / f, weights=f, axis=1)


def poisson_logpmf(x, mu):
    return stats.poisson(mu).logpmf(x)


def poisson_lrt(raw, f, design, refit_mu=True):
    if refit_mu:
        mu_hat_null = poisson_fit_mu_hat(raw, f)
        mu_hat_alt = np.array(
            [poisson_fit_mu_hat(raw[:, design[:, c]], f[:, design[:, c]])
             for c in range(design.shape[1])]
        ).T
    else:
        mu_hat_null = np.mean(raw / f, axis=1)
        mu_hat_alt = np.array(
            [np.mean(raw[:, design[:, c]] / f[:, design[:, c]], axis=1)
             for c in range(design.shape[1])])
    mu_hat_alt_wide = np.dot(mu_hat_alt, design.T)
    null_ll = np.sum(poisson_logpmf(raw, mu_hat_null[:, None] * f), axis=1)
    alt_ll = np.sum(poisson_logpmf(raw, mu_hat_alt_wide * f), axis=1)
    llr = null_ll - alt_ll
    pvalues = stats.chi2(design.shape[1] - 1).sf(-2 * llr)
    return pvalues, llr, mu_hat_null, mu_hat_alt


class Poisson3DeFDR(HiC3DeFDR):
    def estimate_disp(self, estimator='qcml', frac=None, auto_frac_factor=15.,
                      weighted_lowess=True, n_threads=-1):
        # note: all kwargs are ignored
        eprint('estimating dispersion')
        estimator = dispersion.__dict__[estimator] \
            if estimator in dispersion.__dict__ else estimator

        eprint('  loading data')
        disp_idx, _ = self.load_data('disp_idx', 'all')
        row, offsets = self.load_data('row', 'all', idx=disp_idx)
        col, _ = self.load_data('col', 'all', idx=disp_idx)
        scaled, _ = self.load_data('scaled', 'all', idx=disp_idx)

        eprint('  computing pixel-wise mean per condition')
        disp_per_dist = np.zeros((self.dist_thresh_max+1, self.design.shape[1]))
        disp = np.zeros((disp_idx.sum(), self.design.shape[1]))

        def disp_fn(mean):
            return np.zeros_like(mean)

        for c, cond in enumerate(self.design.columns):
            self.save_disp_fn(cond, disp_fn)

        eprint('  saving estimated dispersions to disk')
        self.save_data(disp, 'disp', offsets)
        self.save_data(disp_per_dist, 'disp_per_dist')

    def lrt(self, chrom=None, refit_mu=True, n_threads=-1, verbose=True):
        if chrom is None:
            if n_threads:
                parallel_apply(
                    self.lrt,
                    [{'chrom': c, 'refit_mu': refit_mu, 'verbose': False}
                     for c in self.chroms],
                    n_threads=n_threads
                )
            else:
                for chrom in self.chroms:
                    self.lrt(chrom=chrom, refit_mu=refit_mu)
            return
        eprint('running LRT for chrom %s' % chrom)
        eprint('  loading data', skip=not verbose)
        bias = self.load_bias(chrom)
        size_factors = self.load_data('size_factors', chrom)
        row = self.load_data('row', chrom)
        col = self.load_data('col', chrom)
        raw = self.load_data('raw', chrom)
        disp_idx = self.load_data('disp_idx', chrom)

        eprint('  computing LRT results', skip=not verbose)
        f = bias[row, :][disp_idx, :] * bias[col, :][disp_idx, :] * \
            size_factors[disp_idx, :]
        pvalues, llr, mu_hat_null, mu_hat_alt = poisson_lrt(
            raw[disp_idx, :], f, self.design.values, refit_mu=True)

        if self.loop_patterns:
            eprint('  making loop_idx', skip=not verbose)
            loop_pixels = set().union(
                *sum((load_clusters(pattern.replace('<chrom>', chrom))
                      for pattern in self.loop_patterns.values()), []))
            loop_idx = np.array([True if pixel in loop_pixels else False
                                 for pixel in zip(row[disp_idx],
                                                  col[disp_idx])])
            self.save_data(loop_idx, 'loop_idx', chrom)

        eprint('  saving results to disk', skip=not verbose)
        self.save_data(pvalues, 'pvalues', chrom)
        self.save_data(llr, 'llr', chrom)
        self.save_data(mu_hat_null, 'mu_hat_null', chrom)
        self.save_data(mu_hat_alt, 'mu_hat_alt', chrom)


class Unsmoothed3DeFDR(HiC3DeFDR):
    def estimate_disp(self, estimator='qcml', frac=None, auto_frac_factor=15.,
                      weighted_lowess=True, n_threads=-1):
        # note: all kwargs are ignored
        eprint('estimating dispersion')
        eprint('  loading data')
        disp_idx, _ = self.load_data('disp_idx', 'all')
        row, offsets = self.load_data('row', 'all', idx=disp_idx)
        col, _ = self.load_data('col', 'all', idx=disp_idx)
        scaled, _ = self.load_data('scaled', 'all', idx=disp_idx)

        eprint('  computing pixel-wise mean per condition')
        disp = np.zeros((disp_idx.sum(), self.design.shape[1]))
        for c, cond in enumerate(self.design.columns):
            eprint('  estimating dispersion for condition %s' % cond)
            disp[:, c] = np.maximum(mme_per_pixel(
                scaled[:, self.design[cond]]), 1e-7)

        eprint('  saving estimated dispersions to disk')
        self.save_data(disp, 'disp', offsets)


class Global3DeFDR(HiC3DeFDR):
    def estimate_disp(self, estimator='qcml', frac=None, auto_frac_factor=15.,
                      weighted_lowess=True, n_threads=-1):
        # note: all kwargs except estimator are ignored
        eprint('estimating dispersion')
        estimator = dispersion.__dict__[estimator] \
            if estimator in dispersion.__dict__ else estimator

        eprint('  loading data')
        disp_idx, disp_idx_offsets = self.load_data('disp_idx', 'all')
        loop_idx, _ = self.load_data('loop_idx', 'all')
        row, offsets = self.load_data('row', 'all', idx=disp_idx)
        col, _ = self.load_data('col', 'all', idx=disp_idx)
        raw, _ = self.load_data('raw', 'all', idx=disp_idx)
        f = np.ones_like(raw, dtype=float)
        for i, chrom in enumerate(self.chroms):
            chrom_slice = slice(offsets[i], offsets[i+1])
            row_chrom = row[chrom_slice]
            col_chrom = col[chrom_slice]
            disp_idx_chrom = disp_idx[disp_idx_offsets[i]:disp_idx_offsets[i+1]]
            bias = self.load_bias(chrom)
            size_factors = self.load_data('size_factors', chrom)[disp_idx_chrom]
            f[chrom_slice] = bias[row_chrom, :] * bias[col_chrom, :] \
                * size_factors

        disp = np.zeros((disp_idx.sum(), self.design.shape[1]))
        disp_per_dist = np.zeros((self.dist_thresh_max+1, self.design.shape[1]))
        for c, cond in enumerate(self.design.columns):
            eprint('  estimating dispersion for condition %s' % cond)
            global_disp = estimator(raw[loop_idx, :][:, self.design[cond]],
                                    f=f[loop_idx, :][:, self.design[cond]])
            disp[:, c] = global_disp
            disp_per_dist[:, c] = global_disp

            def disp_fn(mean):
                return np.ones_like(mean) * global_disp

            self.save_disp_fn(cond, disp_fn)

        eprint('  saving estimated dispersions to disk')
        self.save_data(disp, 'disp', offsets)
        self.save_data(disp_per_dist, 'disp_per_dist')
