
# Copyright (C) 2025  Todd Pataky

from math import ceil
import numpy as np
from . metrics import metric_dict
from .. import _spm
from .. _spm import p2string, plist2string, _SPMF
from .. _clusters import ClusterNonparam






class _SnPM(object):
    '''Parent class for all non-parametric SPM classes.'''

    STAT          = 'Z'     #test statistic ("T", "F", "X2" or "T2")
    isanova       = False
    isparametric  = False
    dim           = 0

    def _check_iterations(self, iterations, alpha, force_iterations, nPermTotal):
        if iterations > self.nPermUnique:
            if self.nPermUnique!=-1:
                raise( ValueError('\nNumber of specified iterations (%d) exceeds the maximum possible number of iterations (%d)\n'%(iterations,self.nPermUnique)) )
        elif alpha < (1/nPermTotal):
            raise( ValueError(f'\nSpecified alpha {alpha:0.5f} must be greater than or equal to (1/nPermTotal)={1/nPermTotal:0.5f}\n') )
        elif (((iterations==-1) and (self.nPermUnique>10000)) or (iterations>10000)) and not force_iterations:
            n    = self.nPermUnique if iterations==-1 else iterations
            raise( Warning('\nThe total nuumber of permutations (%d) is very large and may cause computational problems. To enable non-parametric calculations for this many iterations set "force_iterations=True" when calling "inference".\nNOTE: Setting "force_iterations=True" may require substantial computational resources and may cause crashes. USE WITH CAUTION.'%n ))
        elif (iterations!=-1) and (iterations<10):
            raise( ValueError('\nNumber of specified iterations (%d) must be at least 10\n'%iterations) )
        elif (iterations>0) and (iterations<1/alpha):
            raise( ValueError('\nNumber of specified iterations (%d) must be at least %d to conduct inference at alpha=%0.5f\n'%(iterations, ceil(1/alpha), alpha) ))

    def get_nPermUnique_asstr(self):
        n = self.nPermUnique
        if n is None:
            s  = 'None  (child of SnPM list)'
        else:
            s  = str(n) if (n < 1e6) else '%.4g' %n
            s += ' permutations possible'
        return s



'''
################################
   0D SnPM CLASS DEFINITIONS
################################
'''





class _SnPM0D(_SnPM):
    '''Parent class for all 0D non-parametric SPM classes.'''
    def __init__(self, z, mgr, isinlist=False):
        z                   = 0 if np.isnan(z) else z
        self.mgr            = mgr
        self.permuter       = mgr.permuter     #permuter (for conducting inference)
        self.nPermUnique    = None if isinlist else self.permuter.nPermTotal  #number of unique permutations possible
        self.z              = z                #test statistic
        self.isinlist       = isinlist         #is this object in a list (i.e. two- or three-way ANOVA)

    def __repr__(self):
        stat     = self.STAT
        if stat == 'T':
            stat = 't'
        nPermU   = 'inf' if (self.nPermUnique==-1) else str(self.nPermUnique)
        s        = ''
        if self.dim == 0:
            s   += 'SnPM{%s} (0D)\n' %stat
        else:
            s   += 'SnPM{%s}\n' %stat
        if self.isanova:
            s   += '   SPM.effect       :  %s\n'      %self.effect
        s       += '   SPM.z            :  %.3f\n' %self.z
        s       += '   SnPM.nPermUnique :  %s\n'   %self.get_nPermUnique_asstr()
        return s

    def inference(self, alpha=0.05, iterations=-1, force_iterations=False):
        if self.isinlist:
            raise( NotImplementedError( 'Non-parametric inference must be conducted using the parent SnPMList (for two- and three-way ANOVA).' ) )
        self._check_iterations(iterations, alpha, force_iterations, self.permuter.nPermTotal)
        # self.permuter.build_pdf(iterations)
        # zstar     = self.permuter.get_z_critical(alpha)
        # p         = self.permuter.get_p_value(self.z, zstar, alpha)
        self.mgr.permute(iterations)
        zstar    = self.mgr.inference(alpha)
        p        = self.mgr.get_p_value(self.z, zstar, alpha)
        if self.isanova:
            snpm  = SnPM0DiF(self, alpha, zstar, p)
        else:
            snpm  = SnPM0Dinference(self, alpha, zstar, p)
        return snpm




class SnPM0D_T(_SnPM0D):
    STAT = 'T'
    def inference(self, alpha=0.05, two_tailed=True, iterations=-1, force_iterations=False):
        self._check_iterations(iterations, alpha, force_iterations, self.permuter.nPermTotal)
        self.mgr.permute(iterations, two_tailed)
        # alpha0    = 0.5*alpha if two_tailed else alpha
        zstar    = self.mgr.inference(alpha, two_tailed)
        # zstar     = self.permuter.get_z_critical(alpha0, two_tailed)
        p         = self.mgr.get_p_value(self.z, zstar, alpha)
        return SnPM0Dinference(self, alpha, zstar, p, two_tailed)

class SnPM0D_F(_SnPM0D, _SPMF):
    STAT    = 'F'
    isanova = True
    def _repr_summ(self):
        return '{:<5} F = {:<8}\n'.format(self.effect_short,  '%.3f'%self.z)

class SnPM0D_X2(_SnPM0D):
    STAT = 'X2'
class SnPM0D_T2(_SnPM0D):
    STAT = 'T2'





class SnPM0Dinference(_SnPM0D):

    isanova = False

    def __init__(self, spm, alpha, zstar, p, two_tailed=False, isinlist=False):
        super(SnPM0Dinference, self).__init__(spm.z, spm.mgr, isinlist=isinlist)
        self.STAT           = spm.STAT
        self.PDF            = None if isinlist else self.mgr.Z       #permutation PDF
        self.alpha          = alpha                 #Type I error rate
        self.nPerm          = None if isinlist else self.mgr.Z.size  #number of permutations
        self.p              = p                     #P values for each cluster
        self.two_tailed     = two_tailed            #two-tailed test boolean
        self.h0reject       = None                  #null rejection decision
        self.zstar          = zstar                 #critical threshold
        self._check_null()

    def __repr__(self):
        stat     = self.STAT
        if stat == 'T':
            stat = 't'
        s        = ''
        if self.dim == 0:
            s   += 'SnPM{%s} inference (0D)\n' %stat
        else:
            s   += 'SnPM{%s} inference field\n' %stat
        if self.isanova:
            s   += '   SPM.effect         :  %s\n'      %self.effect
        s       += '   SPM.z              :  %.3f\n' %self.z
        s       += '   SnPM.nPermUnique   :  %s\n'   %self.get_nPermUnique_asstr()
        s       += 'Inference:\n'
        if not self.isinlist:
            s   += '   SnPM.nPermActual   :  %d actual permutations\n' %self.nPerm
        s       += '   SPM.alpha          :  %.3f\n' %self.alpha
        if self.two_tailed:
            s   += '   SPM.zstar (lower)  :  %.5f\n' %self.zstar[0]
            s   += '   SPM.zstar (upper)  :  %.5f\n' %self.zstar[1]
        else:
            s   += '   SPM.zstar          :  %.5f\n' %self.zstar
        if self.STAT == 'T':
            s   += '   SPM.two_tailed     :  %s\n'   %str(self.two_tailed)
        s       += '   SPM.h0reject       :  %s\n'   %str(self.h0reject)
        s       += '   SPM.p              :  %s\n' %p2string(self.p)
        return s

    def _check_null(self):
        if self.two_tailed:
            zc0,zc1         = self.zstar
            h0              = (self.z < zc0) or (self.z > zc1)
        else:
            h0              = self.z > self.zstar
        self.h0reject       = h0




class SnPM0DiF(SnPM0Dinference, _SPMF):
    isanova = True
    def _repr_summ(self):
        return '{:<5} F = {:<8} p = {}\n'.format(self.effect.split(' ')[1],  '%.3f'%self.z, p2string(self.p))










'''
################################
   1D SnPM CLASS DEFINITIONS
################################
'''

class _SnPM1D(_SnPM, _spm._SPM):
    '''Parent class for all 1D non-parametric SPM classes.'''

    dim           = 1

    def __init__(self, z, mgr, isinlist=False):
        # z[np.isnan(z)]      = 0
        self.mgr            = mgr              # permutation test manager
        self.permuter       = mgr.permuter     # permuter
        self.Q              = z.size           # field size
        self.nPermUnique    = mgr.permuter.ncomb  # number of unique permutations possible
        self.roi            = mgr.roi          # region(s) of interest
        self.z              = z                # test statistic
        self.isinlist       = isinlist
        self._ClusterClass  = ClusterNonparam

    # def __init__(self, z, perm, isinlist=False):
    #     # z[np.isnan(z)]      = 0
    #     self.permuter       = perm             #permuter (for conducting inference)
    #     self.Q              = z.size           #field size
    #     self.nPermUnique    = perm.nPermTotal  #number of unique permutations possible
    #     self.roi            = perm.roi         #region(s) of interest
    #     self.z              = z                #test statistic
    #     self.isinlist       = isinlist
    #     self._ClusterClass  = ClusterNonparam

    def __repr__(self):
        stat     = self.STAT
        if stat == 'T':
            stat = 't'
        s        = ''
        s       += 'SnPM{%s}\n' %stat
        if self.isanova:
            s   += '   SPM.effect       :  %s\n'      %self.effect
        s       += '   SnPM.z           :  (1x%d) test stat field\n' %self.Q
        s       += '   SnPM.nPermUnique :  %s\n' %self.get_nPermUnique_asstr()
        return s


    def _cluster_inference(self, alpha, clusters, two_tailed=False):
        for cluster in clusters:
            cluster.inference(alpha, self.mgr.Z2, two_tailed)
            if two_tailed:
                cluster.P = min(alpha, 2*cluster.P)
        return clusters

    def _get_clusters(self, zstar, two_tailed, interp, circular, iterations, cluster_metric, z=None):
        clusters      = super(_SnPM1D, self)._get_clusters(zstar, two_tailed, interp, circular, z=z)
        metric        = metric_dict[cluster_metric]
        for c in clusters:
            c.set_metric(metric, iterations, self.nPermUnique, two_tailed)
        return clusters

    # def inference(self, alpha=0.05, iterations=-1, two_tailed=False, interp=True, circular=False, force_iterations=False, cluster_metric='MaxClusterIntegral'):
    #     self._check_iterations(iterations, alpha, force_iterations, self.permuter.nPermTotal)
    #     ### build primary PDF:
    #     self.permuter.build_pdf(iterations)
    #     ### compute critical threshold:
    #     # a          = 0.5*alpha if two_tailed else alpha  #adjusted alpha (if two-tailed)
    #     zstar      = self.permuter.get_z_critical(alpha, two_tailed)
    #     zstar      = zstar[1] if np.size([zstar])==2 else zstar
    #     ### build secondary PDF:
    #     self.permuter.set_metric( cluster_metric )
    #     self.permuter.build_secondary_pdf( zstar, circular )
    #     ### assemble clusters and conduct cluster-level inference:
    #     clusters   = self._get_clusters(zstar, two_tailed, interp, circular, iterations, cluster_metric)
    #     clusters   = self._cluster_inference(alpha, clusters, two_tailed)
    #     if self.isanova:
    #         Fi     = SnPMiF(self, alpha, zstar, clusters)
    #     else:
    #         Fi     = SnPMinference(self, alpha, zstar, two_tailed, clusters)
    #     return Fi

    def inference(self, alpha=0.05, iterations=-1, two_tailed=False, interp=True, circular=False, force_iterations=False, cluster_metric='MaxClusterIntegral'):
        self._check_iterations(iterations, alpha, force_iterations, self.permuter.nPermTotal)
        ### build primary PDF:
        self.mgr.permute( niter=iterations, two_tailed=two_tailed )
        zstar  = self.mgr.inference(alpha, two_tailed=two_tailed)
        # ### build secondary PDF:
        self.mgr.set_metric( cluster_metric )
        self.mgr.build_secondary_pdf( zstar, circular )
        ### assemble clusters and conduct cluster-level inference:
        clusters   = self._get_clusters(zstar, two_tailed, interp, circular, iterations, cluster_metric)
        clusters   = self._cluster_inference(alpha, clusters, two_tailed)
        if self.isanova:
            Fi     = SnPMiF(self, alpha, zstar, clusters)
        else:
            Fi     = SnPMinference(self, alpha, zstar, two_tailed, clusters)
        return Fi
    


        
        # zc = mgr.inference(0.05, two_tailed=True)

    def plot_design(self, **kwdargs):
        msg        = '\n'
        msg       += '"plot_design" is not implemented for non-parametric SPMs. '
        msg       += 'To plot the design matrix use the corresponding parametric procedure and then call "plot_design".\n'
        msg       += 'For example:\n'
        msg       += '   >>>  spm = spm1d.stats.ttest2(yA, yB)\n'
        msg       += '   >>>  spm.plot_design()\n\n'
        raise( NotImplementedError(msg) )


class _SnPM1Donetailed(_SnPM1D):
    def inference(self, alpha=0.05, iterations=-1, interp=True, circular=False, force_iterations=False, cluster_metric='MaxClusterIntegral'):
        return super(_SnPM1Donetailed, self).inference(alpha=alpha, iterations=iterations, two_tailed=False, interp=interp, circular=circular, force_iterations=force_iterations, cluster_metric=cluster_metric)






class SnPM_T(_SnPM1D):
    STAT     = 'T'
class SnPM_F(_SnPM1Donetailed, _SPMF):
    STAT     = 'F'
    isinlist = False
    isanova  = True
    def _repr_summ(self):
        return '{:<5} F = {:<18}\n'.format(self.effect_short,  self._repr_teststat_short())
class SnPM_X2(_SnPM1Donetailed):
    STAT = 'X2'
class SnPM_T2(_SnPM1Donetailed):
    STAT = 'T2'





class SnPMinference(_SnPM1D, _spm._SPMinference):
    def __init__(self, spm, alpha, zstar, two_tailed, clusters):
        super(SnPMinference, self).__init__(spm.z, spm.mgr)
        self.STAT           = spm.STAT
        self.PDF0           = self.mgr.Z        #primary permutation PDF
        self.PDF1           = self.mgr.Z2       #secondary PDF (cluster-level)
        self.alpha          = alpha               #Type I error rate
        self.clusters       = clusters            #supra-threshold cluster information
        self.nClusters      = len(clusters)         #number of supra-threshold clusters
        self.p              = [c.P for c in clusters]  #P values for each cluster
        self.two_tailed     = two_tailed          #two-tailed test boolean
        self.zstar          = zstar               #critical threshold
        self.h0reject       = self.nClusters > 0  #null hypothesis rejection decision
        self.roi            = spm.roi



    def __repr__(self):
        s        = ''
        s       += 'SnPM{%s} inference field\n' %self.STAT
        if self.isanova:
            s   += '   SPM.effect         :  %s\n'      %self.effect
        s       += '   SPM.z              :  (1x%d) raw test stat field\n' %self.Q
        s       += '   SnPM.nPermUnique   :  %s\n'      %self.get_nPermUnique_asstr()
        s       += 'Inference:\n'
        # s       += '   SnPM.nPermActual   :  (%d) actual permutations\n' %self.nPerm0
        s       += '   SPM.alpha          :  %.3f\n'    %self.alpha
        s       += '   SPM.zstar          :  %.5f\n'    %self.zstar
        s       += '   SPM.h0reject       :  %s\n'      %self.h0reject
        # s       += '   SPM.cluster_metric :  %s\n'   %self.cluster_metric
        s       += '   SPM.p              :  (%s)\n'   %plist2string(self.p)
        return s



class SnPMiF(SnPMinference, _SPMF):
    def __init__(self, spm, alpha, zstar, clusters):
        super(SnPMiF, self).__init__(spm, alpha, zstar, False, clusters)
        self.isanova        = True
        self.effect         = spm.effect
        self.effect_short   = spm.effect_short

    def _repr_summ(self):
        return '{:<5} z={:<18} h0reject={}\n'.format(self.effect_short,  self._repr_teststat_short(), self.h0reject)





