"""
Created on Sep 24, 2013
@author: casale
"""
import scipy as SP
import scipy.linalg as LA
import pdb

import sys
sys.path.append('./../../..')
import mtSet.pycore.external.limix.plink_reader as plink_reader


def genBinormal(dim1,dim2,percSign=0.5,std=1e-1):
    rv = (2*(SP.rand(dim1,dim2)>percSign)-1)+std*SP.randn(dim1,dim2)
    return rv

def selectRnd(n_sel,n_all):
    rv = SP.zeros(n_all)
    rv[:n_sel] = 1
    SP.random.shuffle(rv)
    rv = rv==1
    return rv

class CSimulator:
    """
    this class takes care of phenotype generation in a flexible way
    """

    def __init__(self,bfile,XX=None,P=1):
        """
        X: genotype matrix
        traitNum: number of traits to be considered
        """
        self.bfile = bfile
        self.N = XX.shape[0]
        self.P = P
        self.XX = XX

        pass
    
    def getRegion(self,size=3e4,min_nSNPs=1,chrom_i=None,pos_min=None,pos_max=None):
        """
        Sample a region from the piece of genotype X, chrom, pos
        minSNPnum:  minimum number of SNPs contained in the region
        Ichrom:  restrict X to chromosome Ichrom before taking the region
        cis:        bool vector that marks the sorted region
        region:  vector that contains chrom and init and final position of the region
        """
        bim = plink_reader.readBIM(self.bfile,usecols=(0,1,2,3))
        
        chrom = SP.array(bim[:,0],dtype=int)
        pos   = SP.array(bim[:,3],dtype=int)
        
        if chrom_i is None:
            n_chroms = chrom.max()
            chrom_i  = int(SP.ceil(SP.rand()*n_chroms))

        pos   = pos[chrom==chrom_i]
        chrom = chrom[chrom==chrom_i]

        ipos = SP.ones(len(pos),dtype=bool)
        if pos_min is not None:
            ipos = SP.logical_and(ipos,pos_min<pos)

        if pos_max is not None:
            ipos = SP.logical_and(ipos,pos<pos_max)

        pos = pos[ipos]
        chrom = chrom[ipos]
        
        if size==1:
            # select single SNP
            idx = int(SP.ceil(pos.shape[0]*SP.rand()))
            cis  = SP.arange(pos.shape[0])==idx
            region = SP.array([chrom_i,pos[idx],pos[idx]])
        else:
            while 1:
                idx = int(SP.floor(pos.shape[0]*SP.rand()))
                posT1 = pos[idx]
                posT2 = pos[idx]+size
                if posT2<=pos.max():
                    cis = chrom==chrom_i
                    cis*= (pos>posT1)*(pos<posT2)
                    if cis.sum()>min_nSNPs: break
            region = SP.array([chrom_i,posT1,posT2])

        start = SP.nonzero(cis)[0].min()
        nSNPs  = cis.sum()
        rv = plink_reader.readBED(self.bfile,useMAFencoding=True,start = start, nSNPs = nSNPs,bim=bim)
        Xr = rv['snps']
        return Xr, region


    def genRegionTerm(self,X,vTot=0.1,pCausal=0.10,nCausal=None,pCommon=1.,nCommon=None,plot=False,distribution='biNormal'):
        """
        Generate population structure term
        Population structure is simulated by background SNPs
        
        beta_pdf:        pdf used to generate the regression weights
                          for now either Normal or fixed
        variance:        variance of the term
        percCausal:    percentage of causal SNPs
        Xcausal:          set of SNPs being causal
        """
        S = X.shape[1]

        # number of causal, common, specific
        if nCausal==None:
            nCausal=int(SP.floor(pCausal*S))
        if nCommon==None:
            nCommon = round(pCommon*nCausal)
        nSpecific = self.P*(nCausal-nCommon)

        # common SNPs
        if nCommon>0:
            if distribution=='biNormal':
                Bc  = SP.kron(genBinormal(nCommon,1),genBinormal(1,self.P))
            elif distribution=='normal':
                Bc  = SP.kron(SP.randn(nCommon,1),SP.randn(1,self.P))
            Ic  = selectRnd(nCommon,S)
            Yc  = SP.dot(X[:,Ic],Bc)
            Yc *= SP.sqrt(nCommon/Yc.var(0).mean())