#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  phase_transition_hmm.py
#
#  Copyright 2018 scimmia <scimmia@scimmia-ThinkPad-L540>
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#
#

import json
import os
import pathlib
import sys
import pickle
import string
import random

from sklearn.metrics import roc_auc_score
from scipy.stats import ranksums

from b2bTools.singleSeq.PSPer.hmmer_research.hmmsearch_otf import hmmer_wrapper
from b2bTools.singleSeq.PSPer.source.utils import seqs_features,leggifasta
import numpy as np

if sys.version_info[1] <= 12:
    # Python version is less than 3.12, so use importlib to import the library
    from pomegranate import DiscreteDistribution
    from pomegranate import IndependentComponentsDistribution
    from pomegranate import State
    from pomegranate import HiddenMarkovModel
else:
    # Python version is 3.13 or higher, so import the library directly
    from pomegranate.distributions.categorical import Categorical as DiscreteDistribution
    from pomegranate.distributions.independent_components import IndependentComponents as IndependentComponentsDistribution
    from pomegranate.hmm import SparseHMM as HiddenMarkovModel

from scipy.optimize import minimize

from b2bTools.singleSeq.DynaMine.Predictor import DynaMine
from b2bTools.singleSeq.EFoldMine.Predictor import EFoldMine
from b2bTools.singleSeq.DisoMine.Predictor import DisoMine

def load_pickle(pickle_file):
    try:
        with open(pickle_file, 'rb') as f:
            pickle_data = pickle.load(f)
    except UnicodeDecodeError as e:
        with open(pickle_file, 'rb') as f:
            pickle_data = pickle.load(f, encoding='latin1')
    except Exception as e:
        print('Unable to load data ', pickle_file, ':', e)
        raise
    return pickle_data

class phase_hmm:
    def __init__(self):
        self.emiss_diverse=None
        self.emiss_bg=None
        self.scaler=None

        self.scriptDir = str(pathlib.Path(__file__).parent.absolute())

    def fit(self):
        if 'emissions.m' in os.listdir(os.path.join(self.scriptDir,'marshalled')):
            self.emiss_diverse,self.emiss_bg=load_pickle(os.path.join(self.scriptDir,'marshalled/emissions.m'))
        else:
            self.emiss_diverse,self.emiss_bg=self.learn_distro()
            pickle.dump((self.emiss_diverse,self.emiss_bg),open(os.path.join(self.scriptDir,'marshalled/emissions.m','w')))

        self.build_hmm()

        if 'scaler_new.m' in os.listdir(os.path.join(self.scriptDir,'marshalled')):
            self.scaler=load_pickle(os.path.join(self.scriptDir,'marshalled/scaler_new.m'))
        else:
            self.scaler=self.learn_scaling()
            pickle.dump(self.scaler,open(os.path.join(self.scriptDir,'marshalled/scaler_new.m','w')))


    def build_hmm(self,param=None,optim_params=False,PLOT_DISTROS=False):
        if self.emiss_diverse==None or self.emiss_bg==None:
            raise ValueError('YOU FIRST NEED TO FIT EMISSIONS! RUN .fit()')

        diverso=self.emiss_diverse
        costante=self.emiss_bg

        if not optim_params:
            len_pld=80
            len_rna=50
            len_altro=20
            len_grigio=15

            ### param transizioni ###
            TRANS_ALTRO=0.001

            ### somma a 1 ###
            TRANS_ESTENDI_STATO=0.6
        else:
            len_pld=int(param[1])
            len_rna=int(param[2])
            len_altro=int(param[3])
            len_grigio=int(param[4])

        diz_rna={}
        diz_NON_rna={}

        for i in range(11):
            if i == 10:
                diz_rna['a'] = 0.05 + i * 0.1 * 0.9 /5.5
            else:
                diz_rna[str(i)] = 0.05 + i * 0.1 * 0.9
        for i in range(10, -1, -1):
            if i != 10:
                diz_NON_rna[str(i)] = 0.05 + (10-i) * 0.1 * 0.9
            else:
                diz_NON_rna['a'] = 0.05 + (10 - i) * 0.1 * 0.9

        norm=sum(diz_rna.values())

        for i in diz_rna.keys():
            diz_rna[i]=diz_rna[i]/norm

        norm=sum(diz_NON_rna.values())
        for i in diz_NON_rna.keys():
            diz_NON_rna[i]=diz_NON_rna[i]/norm

        rna_alto=DiscreteDistribution(diz_rna)
        rna_basso=DiscreteDistribution(diz_NON_rna)

        distrib_altro = IndependentComponentsDistribution(
            [
                DiscreteDistribution(costante[0]),
                DiscreteDistribution(costante[1]),
                DiscreteDistribution(costante[2]),
                rna_basso,
                DiscreteDistribution(costante[4])
            ]
        )

        distrib_grigia = IndependentComponentsDistribution(
            [
                DiscreteDistribution(costante[0]),
                DiscreteDistribution(diverso[1]),
                DiscreteDistribution(costante[2]),
                rna_basso,
                DiscreteDistribution(diverso[4])
            ]
        )

        distrib_rna = IndependentComponentsDistribution(
            [
                DiscreteDistribution(costante[0]),
                DiscreteDistribution(costante[1]),
                DiscreteDistribution(costante[2]),
                rna_alto,
                DiscreteDistribution(costante[4])
            ]
        )
        distrib_pdl = IndependentComponentsDistribution(
            [
                DiscreteDistribution(diverso[0]),
                DiscreteDistribution(costante[1]),
                DiscreteDistribution(diverso[2]),
                rna_basso,
                DiscreteDistribution(diverso[4])
            ]
        )

        PLD1=[]
        PLD2=[]

        altroB=[]
        altroF=[]

        altroRNA=[]
        altroPLD=[]

        grigioB=[]
        grigioF=[]

        grigio_intra1=[]
        grigio_intra2=[]

        grigioRNA=[]
        grigioPLD=[]

        RNA1=[]
        RNA2=[]

        for i in range(len_pld):
            PLD1+=[State(distrib_pdl, name='PLD_1'+str(i))]
            PLD2+=[State(distrib_pdl, name='PLD_2'+str(i))]

        for i in range(len_rna):
            RNA1+=[State(distrib_rna, name='RRM_1'+str(i))]
            RNA2+=[State(distrib_rna, name='RRM_2'+str(i))]

        for i in range(len_altro):
            altroB+=[State(distrib_altro, name='OTHER_B_'+str(i))]
            altroPLD+=[State(distrib_altro, name='OTHER_PLD_'+str(i))]
            altroRNA=[State(distrib_altro, name='OTHER_RNA_'+str(i))]
            altroF+=[State(distrib_altro, name='OTHER_F_'+str(i))]

        for i in range(len_grigio):
            grigioB+=[State(distrib_grigia, name='SPACER_B_'+str(i))]
            grigioF+=[State(distrib_grigia, name='SPACER_F_'+str(i))]
            grigioPLD+=[State(distrib_grigia, name='SPACER_PLD_'+str(i))]
            grigioRNA+=[State(distrib_grigia, name='SPACER_RNA_'+str(i))]
            grigio_intra1+=[State(distrib_grigia, name='SPACER_intra2_'+str(i))]
            grigio_intra2+=[State(distrib_grigia, name='SPACER_intra1_'+str(i))]

        model = HiddenMarkovModel()
        background = HiddenMarkovModel()

        bg = State(distrib_altro,name='background')

        background.add_states(bg)
        background.add_transition(bg ,bg ,0.7)
        background.add_transition(bg ,background.end ,0.3)
        background.add_transition(background.start,bg,1)
        background.bake()

        self.background=background

        # grigio
        for i in range(len(grigioB)):
            model.add_states(grigioB[i])
        for i in range(len(grigioF)):
            model.add_states(grigioF[i])
        for i in range(len(grigioRNA)):
            model.add_states(grigioRNA[i])
        for i in range(len(grigioPLD)):
            model.add_states(grigioPLD[i])
        for i in range(len(grigio_intra1)):
            model.add_states(grigio_intra1[i])
        for i in range(len(grigio_intra2)):
            model.add_states(grigio_intra2[i])

        # altro
        for i in range(len(altroB)):
            model.add_states(altroB[i])
        for i in range(len(altroF)):
            model.add_states(altroF[i])
        for i in range(len(altroPLD)):
            model.add_states(altroPLD[i])
        for i in range(len(altroRNA)):
            model.add_states(altroRNA[i])

        #######################################
        #PLD
        for i in range(len(PLD1)):
            model.add_states(PLD1[i])
        for i in range(len(PLD2)):
            model.add_states(PLD2[i])

        #RNA
        for i in range(len(RNA1)):
            model.add_states(RNA1[i])
        for i in range(len(RNA2)):
            model.add_states(RNA2[i])

        #######################################
        #              GRIGIO
        for i in range(len(grigioB)-1):
            model.add_transition(grigioB[i],grigioB[i+1] , 1)
        model.add_transition(grigioB[-1],grigioB[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(grigioF)-1):
            model.add_transition(grigioF[i],grigioF[i+1] , 1)
        model.add_transition(grigioF[-1],grigioF[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(grigioPLD)-1):
            model.add_transition(grigioPLD[i],grigioPLD[i+1] , 1)
        model.add_transition(grigioPLD[-1],grigioPLD[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(grigioRNA)-1):
            model.add_transition(grigioRNA[i],grigioRNA[i+1] , 1)
        model.add_transition(grigioRNA[-1],grigioRNA[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(grigio_intra1)-1):
            model.add_transition(grigio_intra1[i],grigio_intra1[i+1] , 1)
        model.add_transition(grigio_intra1[-1],grigio_intra1[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(grigio_intra2)-1):
            model.add_transition(grigio_intra2[i],grigio_intra2[i+1] , 1)
        model.add_transition(grigio_intra2[-1],grigio_intra2[-1] , TRANS_ESTENDI_STATO)

        ################################################
        #                 altro
        for i in range(len(altroB)-1):
            model.add_transition(altroB[i],altroB[i+1] , 1)
        model.add_transition(altroB[-1],altroB[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(altroPLD)-1):
            model.add_transition(altroPLD[i],altroPLD[i+1] , 1)
        model.add_transition(altroPLD[-1],altroPLD[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(altroF)-1):
            model.add_transition(altroF[i],altroF[i+1] , 1)
        model.add_transition(altroF[-1],altroF[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(altroRNA)-1):
            model.add_transition(altroRNA[i],altroRNA[i+1] , 1)
        model.add_transition(altroRNA[-1],altroRNA[-1] , TRANS_ESTENDI_STATO)

        ################################################
        #                  rna
        for i in range(len(RNA1)-1):
            model.add_transition(RNA1[i],RNA1[i+1] , 1)
        model.add_transition(RNA1[-1],RNA1[-1] , TRANS_ESTENDI_STATO)
        for i in range(len(RNA2)-1):
            model.add_transition(RNA2[i],RNA2[i+1] , 1)
        model.add_transition(RNA2[-1],RNA2[-1] , TRANS_ESTENDI_STATO)

        ################################################
        #                 pld
        for i in range(len(PLD1)-1):
            model.add_transition(PLD1[i],PLD1[i+1] , 1)
        model.add_transition(PLD1[-1],PLD1[-1] , TRANS_ESTENDI_STATO)

        for i in range(len(PLD2)-1):
            model.add_transition(PLD2[i],PLD2[i+1] , 1)
        model.add_transition(PLD2[-1],PLD2[-1] , TRANS_ESTENDI_STATO)

        ##################### MODULO 1 #######################
        #  somma a 0.05 per avere estensione di 20 al blocco #
        # la P disponibile e' 1-P(self-transizione interna al blocco) --> ora a 0.95 --> estensione media 20
        ######################################################


        ######################################################à
        #                    BEGIN                            #
        #######################################################

        model.add_transition(model.start,PLD1[0] , 0.5)
        model.add_transition(model.start,grigioB[0] , 0.3)
        model.add_transition(model.start,altroB[0] , TRANS_ALTRO)
        model.add_transition(model.start,RNA1[0] , 0.1)
        ######################################################à
        #                    altrob                           #
        #######################################################
        model.add_transition(altroB[-1],PLD1[0] , 0.1*(1-TRANS_ESTENDI_STATO))
        model.add_transition(altroB[-1],grigioB[0] , 0.8*(1-TRANS_ESTENDI_STATO))
        model.add_transition(altroB[-1],RNA1[0] , 0.1*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    grigioB                            #
        #######################################################
        model.add_transition(grigioB[-1],PLD1[0] , 0.45*(1-TRANS_ESTENDI_STATO))
        model.add_transition(grigioB[-1],altroB[0] , TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))
        model.add_transition(grigioB[-1],RNA1[0] , 0.45*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    RNA1                            #
        #######################################################
        model.add_transition(RNA1[-1],PLD2[0] , 0.1*(1-TRANS_ESTENDI_STATO))
        model.add_transition(RNA1[-1],grigio_intra1[0] , 0.6*(1-TRANS_ESTENDI_STATO))
        model.add_transition(RNA1[-1],grigioRNA[0] , 0.4*(1-TRANS_ESTENDI_STATO))
        model.add_transition(RNA1[-1],altroRNA[0] , TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))

        model.add_transition(grigio_intra1[-1],RNA1[0] , 1*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    grigioRNA                        #
        #######################################################
        model.add_transition(grigioRNA[-1],PLD2[0] , 0.8*(1-TRANS_ESTENDI_STATO))
        model.add_transition(grigioRNA[-1],altroRNA[0] , TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))


        ######################################################à
        #                    altroRNA                        #
        #######################################################
        model.add_transition(altroRNA[-1],PLD2[0] , 0.1*(1-TRANS_ESTENDI_STATO))
        model.add_transition(altroRNA[-1],grigioRNA[0] , 0.9*(1-TRANS_ESTENDI_STATO))


        ######################################################à
        #                   PLD2                             #
        #######################################################
        model.add_transition(PLD2[-1],altroF[0] ,TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))
        model.add_transition(PLD2[-1],grigioF[0] , 0.1*(1-TRANS_ESTENDI_STATO))
        model.add_transition(PLD2[-1],model.end ,0.8*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    altroF                      #
        #######################################################
        model.add_transition(altroF[-1],model.end , 0.2*(1-TRANS_ESTENDI_STATO))
        model.add_transition(altroF[-1],grigioF[0] , 0.8*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    grigioF                      #
        #######################################################
        model.add_transition(grigioF[-1],model.end , 0.2*(1-TRANS_ESTENDI_STATO))
        model.add_transition(grigioF[-1],altroF[0] , 0.8*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                   PLD1                             #
        #######################################################
        model.add_transition(PLD1[-1],altroPLD[0] , TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))
        model.add_transition(PLD1[-1],grigioPLD[0] , 0.8*(1-TRANS_ESTENDI_STATO))
        model.add_transition(PLD1[-1],RNA2[0] ,0.1*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    grigioPLD                        #
        #######################################################
        model.add_transition(grigioPLD[-1],RNA2[0] , 0.8*(1-TRANS_ESTENDI_STATO))
        model.add_transition(grigioPLD[-1],altroPLD[0] , TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))


        ######################################################à
        #                    altroPLD                       #
        #######################################################
        model.add_transition(altroPLD[-1],RNA2[0] , 0.2*(1-TRANS_ESTENDI_STATO))
        model.add_transition(altroPLD[-1],grigioPLD[0] , 0.8*(1-TRANS_ESTENDI_STATO))

        ######################################################à
        #                    RNA2                            #
        #######################################################
        model.add_transition(RNA2[-1],grigio_intra2[0] , 0.4*(1-TRANS_ESTENDI_STATO))
        model.add_transition(RNA2[-1],grigioF[0] , 0.4*(1-TRANS_ESTENDI_STATO))
        model.add_transition(RNA2[-1],altroF[0] , TRANS_ALTRO*(1-TRANS_ESTENDI_STATO))
        model.add_transition(RNA2[-1],model.end , 0.2*(1-TRANS_ESTENDI_STATO))
        model.add_transition(grigio_intra2[-1],RNA2[0] , 1*(1-TRANS_ESTENDI_STATO))
        ######

        model.bake()

        self.model=model

    def build_vector(self, seqs):
        SEQS = []
        for key, value in seqs.items():
            temp = (key, value)
            SEQS.append(temp)


        # TODO - could still improve this by providing the predictions if already run prior to PSPer request! Will speed things up...
        dm = DynaMine(predictionTypes=['backbone', 'sidechain'])
        dm.predictSeqs(SEQS)
        efm = EFoldMine()
        efm.allPredictions = dm.allPredictions
        efm.predictSeqs(SEQS)
        de = DisoMine()
        de.allPredictions = efm.allPredictions
        disomine=de.predictSeqs(SEQS)
        entry={}
        for protID in range(len(seqs)):
            entry[json.loads(de.getAllPredictionsJson('Disomine'))['results'][protID]['proteinID']] = np.array(json.loads(de.getAllPredictionsJson('Disomine'))['results'][protID]['disoMine'])

        disomine = entry
        hmmer_wrapper_instance = hmmer_wrapper(root=os.path.join(self.scriptDir, 'hmmer_research/'))
        hmmer_results = hmmer_wrapper_instance.predict(seqs)

        hmmer_features = seqs_features(seqs)
        result_dict = {}

        for hmmer_key in hmmer_results.keys():
            current_hmmer_result = hmmer_results[hmmer_key]
            current_features = hmmer_features[hmmer_key]
            current_disorder = disomine[hmmer_key]

            result_dict[hmmer_key] = discretizza([
                current_features[0],
                current_features[1],
                current_features[2],
                current_hmmer_result,
                current_disorder
            ])

        return result_dict

    def predict_proba(self,seq,scale=True):

        if type(seq)==str:
            seq={'input_seq':seq}
            vets=self.build_vector(seq)
        elif type(seq)==dict and type(seq[list(seq.keys())[0]])==str:
            vets=self.build_vector(seq)
        elif type(seq)==dict and type(seq[list(seq.keys())[0]])==list:
            vets = seq
        else:
            raise TypeError('predict takes a protein sequnce (str) or a dict of sequences or a dict of feature vectors (list)')
        res={}

        for i in vets.keys():
            true=self.model.log_probability(vets[i])
            rand=self.background.log_probability(vets[i])
            score=true-rand

            if scale==True:
                cattivo=False

                if score<-10000:
                    score=-10000
                    cattivo=True
                score=self.scaler.transform([[score]])[0][0]
                if cattivo:
                    score=0

                if score<0:
                    score=0.0
                elif score>1:
                    score=1.0
            res[i]=float(score)
        return res

    def viterbi(self,seq,scale=False):
        if type(seq)==str:
            seq={'input_seq':seq}
            vet=self.build_vector(seq)
        elif type(seq)==dict and type(seq[list(seq.keys())[0]])==str:
            vet=self.build_vector(seq)
        elif type(seq)==dict and type(seq[list(seq.keys())[0]])==list:
            vet = seq
        else:
            raise TypeError('predict takes a protein sequnce (str) or a dict of sequences or a dict of feature vectors (list)')
        r={}
        for i in seq.keys():
            p=[]
            seq_key = vet[i]
            viterbi_model = self.model.viterbi(seq_key)[1]
            if not viterbi_model:
                # print(f"No Viterbi model for input = {seq_key}")
                continue
            else:
                path= ", ".join(state.name for i, state in viterbi_model)
                for k in path.split(',')[1:-1]:
                    p+=[k.split('_')[0]]

                assert len(p)==len(seq[i])

                r[i]=p

        return r

    def learn_scaling(self):
        from sklearn.preprocessing import MinMaxScaler

        if FAST and 'backgroundDataset.fasta.cPickle' in os.listdir('marshalled'):
            vets_bg=load_pickle('marshalled/backgroundDataset.fasta.cPickle')

        else:
            print ('no vect marshal for background, this will take ages')
            seqs_bg=leggifasta('validationDatasets/background/backgroundDataset.fasta')
            s={}

            for i in seqs_bg.keys()[:]:
                if len(seqs_bg[i])>=MAX_SEQ_SIZE or len(seqs_bg[i])<MIN_SEQ_SIZE:
                    continue
                s[i]=seqs_bg[i]
            seqs_bg=s
            vets_bg=self.build_vector(seqs_bg)
            pickle.dump((vets_bg,seqs_bg),open('marshalled/backgroundDataset.fasta.cPickle','w'))

        if "phaseSepVects.cPickle" in os.listdir('marshalled'):
            vets=load_pickle('marshalled/phaseSepVects.cPickle')
        else:
            seqs=leggifasta('validationDatasets/phaseSepProts/phase_separation_proteins.fasta')
            print ('no vect marshal for phase_separation, marshalling.',len(seqs),'proteins to do')
            s={}
            for i in seqs.keys()[:]:

                if len(seqs[i])>=MAX_SEQ_SIZE:
                    continue
                s[i]=seqs[i]
            seqs=s
            vets=self.build_vector(seqs)
            pickle.dump(vets,open('marshalled/phaseSepVects.cPickle','w'))


        if "RNABPdatasetFinal.fasta.cPickle" in os.listdir('marshalled'):
            vetsbp=load_pickle('marshalled/RNABPdatasetFinal.fasta.cPickle')
            #del vetsbp['Q9UQ35']
            #pickle.dump(vetsbp,open('marshalled/RNABPdatasetFinal.fasta.cPickle','w'))
        else:
            seqsbp=leggifasta('validationDatasets/RNAbp/RNABPdatasetFinal.fasta')
            print ('no vect marshal for phase_separation, marshalling.',len(seqsbp),'proteins to do')
            s={}
            for i in seqsbp.keys()[:]:

                if len(seqsbp[i])>=MAX_SEQ_SIZE:
                    continue
                s[i]=seqsbp[i]
            seqsbp=s
            vetsbp=self.build_vector(seqsbp)


            pickle.dump(vetsbp,open('marshalled/RNABPdatasetFinal.fasta.cPickle','w'))
        #vets.update(vets_bg)
        #print vetsbp[0]
        #print vets
        vetsbp.update(vets_bg)
        if vetsbp.has_key('Q25434'):
            del vetsbp['Q25434']
        r=self.predict_proba(vetsbp,scale=False)
        scaler=MinMaxScaler() #QuantileTransformer()
        scaler.fit(np.array(r.values()).reshape(-1, 1))
        a=[]
        pickle.dump(scaler, open('marshalled/scaler_new.m', 'w'))
        return scaler

    def learn_distro(self,bg_file='validationDatasets/background/backgroundDataset.fasta',bins=[12,12,12,10,10],reverse=[False,True,True,True,True],minemiss=0.05,plot_histograms=True): # si e' grezza,ma in bin_range_per_fea devi mettere, per ogni feature il numero di bins e il range di ogni feature. questo perche' alcuni valori non vengono visti mai
        seqs=leggifasta(bg_file)
        s={}
        for i in seqs.keys()[:500]:

            if len(seqs[i])>MAX_SEQ_SIZE or len(seqs[i])<MIN_SEQ_SIZE:
                continue
            s[i]=seqs[i]
        seqs=s

        if FAST and 'seqs_vec_background.m' in os.listdir('marshalled'):
            vets,seqs=load_pickle('marshalled/seqs_vec_background.m')
        else:
            vets=self.build_vector(seqs)
        printable=string.printable

        cont=0
        g=[]
        for i in vets.keys():
            dig=vets[i]
            for j in dig:
                l=[]
                for k in j:
                    l+=[printable.index(k)]
                g+=[l]
        diz_prob=[]
        diz_cost=[]
        g=np.array(g)
        h=[]
        for i in range(g.shape[1]):
            hist, bin_edges=np.histogram(g[:,i], bins=bins[i], range=(0,bins[i]), density=True,normed=True)
            h+=[hist]
            costante=list(np.cumsum(hist))

            for k in range(len(costante)):

                costante[k]=max(costante[k],minemiss)
                #print costante[i],max(costante[i],minemiss)
                #raw_input()
            #print min(costante),'qe'
            assert min(costante)>=minemiss
            d=[]
            if reverse[i]:
                prev=1.0
                pc=1.0
                c=[1]
                for k in range(1,len(hist[1:])+1):
                    c+=[max(c[-1]-hist[k-1],minemiss)]
                    assert min(c)>=0.05
                for k in range(len(hist)-1,-1,-1):
                    prev-=hist[k]
                    prev=max(prev,minemiss)
                    d=[prev]+d
                diz_cost+=[c]
                diz_prob+=[d]
            else:
                prev=1.0
                diz_cost+=[costante]

                for k in range(len(hist)):
                    prev-=hist[k]
                    prev=max(prev,minemiss)
                    d+=[prev]
                diz_prob+=[d]

        distro=[]
        for i in diz_prob:
            diz={}
            for k in range(len(i)):
                diz[printable[k]]=i[k]
            distro+=[diz]

        distro_cost=[]
        cont=0
        for i in diz_cost:

            di={}
            for k in range(len(i)):
                #print i[k]
                assert i[k]>=0.05
                di[printable[k]]=i[k]
            distro_cost+=[di]
            cont+=1
        # if plot_histograms:
        # 	fea_names=['Sequence complexity','Arg enrichment','Tyr enrichment','RNA-binding','Disorder']
        # 	for i in range(g.shape[1]):
        # 		data = g[:,i]
        # 		fig, ax1 = plt.subplots()
        # 		d = np.diff(np.unique(data)).min()
        # 		left_of_first_bin = data.min() - float(d)/2
        # 		right_of_last_bin = data.max() + float(d)/2
        # 		ax1.hist(data, np.arange(left_of_first_bin, right_of_last_bin + d, d),normed=True,label='feature distribution')

        # 		#ax1.hist(scores, bins=15,alpha=0.3, color="b")
        # 		#ax1.set_xlabel('HMM scores')
        # 		# Make the y-axis label, ticks and tick labels match the line color.
        # 		#ax1.set_ylabel('Occurrence', color='b')
        # 		ax1.tick_params('y', colors='b')

        # 		ax2 = ax1.twinx()
        # 		ax2.tick_params('y', colors='r')
        # 		#values, base = np.histogram(scores, bins=100,normed=False)

        # 		y=[]
        # 		cont=0.0
        # 		x=[]
        # 		for j in sorted(distro_cost[i].keys()):
        # 			if distro_cost[i][j]==0.05:
        # 				y+=[0.0]
        # 			elif distro_cost[i][j]==0.95:
        # 				y+=[1]
        # 			else:
        # 				y+=[distro_cost[i][j]]
        # 			x+=[cont]
        # 			cont+=1

        # 		if i!=0:
        # 			yc=[]
        # 			for k in y:
        # 				yc+=[1-k]
        # 				y=yc

        # 		ax2.plot(x, y, c='r',label='Cumulative')
        # 		if i==1:
        # 			plt.legend()

        # 		#plt.plot(x,y,label='feature cumulative distribution', alpha=0.4)
        # 		plt.legend()

        # 		plt.ylabel('Probability')
        # 		plt.xlabel('Feature value')

        # 		plt.title(fea_names[i])
        # 		plt.savefig(fea_names[i]+'.png',dpi=400)
        # 		plt.clf()
        for i in range(len(distro_cost)):
            v=distro_cost[i].values()
            for k in distro_cost[i].keys():
                distro_cost[i][k]=distro_cost[i][k]/sum(v)
        for i in range(len(distro)):
            v=distro[i].values()
            for k in distro[i].keys():
                distro[i][k]=distro[i][k]/sum(v)
        pickle.dump((distro,distro_cost), open('marshalled/emissions.m', 'w'))
        pickle.dump((distro,distro_cost), open('marshalled/emissions.m', 'w'))
        return distro,distro_cost

class scoring_function:
    def __init__(self):
        self.distro,self.distro_cost=load_pickle('marshalled/emissions.m')
        seqs_rna=leggifasta('rna_bp.fasta')
        s={}
        for i in seqs_rna.keys()[:]:
            s[i]=seqs_rna[i]
        seqs_rna=s
        seqs=leggifasta('phase_separation_proteins.fasta')

        s={}
        for i in seqs.keys()[:]:
            s[i]=seqs[i]
        seqs=s
        self.vets=build_vector(seqs)
        self.vets_rna=build_vector(seqs_rna)

    def score(self,c):
        model=phase_hmm(self.distro,self.distro_cost,param=c,optim_params=True)
        a= model.predict(self.vets)
        b= model.predict(self.vets_rna)
        return np.median(b)-np.median(a)

def parameter_opt():
    score_obj=scoring_function()
    c=minimize(score_obj.score, [0,70,80,50,30])
    print (c)

def discretizza(feas,continue_val=[False,False,False,False,True],n_bins=10): ## metti numero bins per ogni feature
    printable=string.printable
    bins = np.linspace(0, 1, n_bins)
    dig=[]
    for i in range(len(continue_val)):
        if continue_val[i]:
            dig+=[np.digitize(feas[i], bins)]
            #print np.digitize(feas[i], bins)

        else:

            dig+=[feas[i]]
    v=[]
    for i in range(len(dig[0])):
        s=''
        for j in range(len(dig)):
            #print len(feas[j]),len(feas),len(threshold),i,j
            s+=printable[dig[j][i]]
        v+=[s]
    return v

####  cose da fare : cumulativa background direttamente da hist, le distro fatte  bene--> diverso[i], costante[i] per fare background o non background
PSP = 'validationDatasets/phaseSepProts/phase_separation_proteins.fasta'
FAST = True
#TARGET = "validationDatasets/mip6.fasta"
#TARGET = "proteomes/ecoli.UP000000625.fasta"
#TARGET = "proteomes/saccharomycesUP000002311.fasta"
#TARGET= "validationDatasets/144granuleCoreProts/144granuleCoreProts.fasta"
#TARGET = "validationDatasets/granuleForming/granuleForming.fasta"
#TARGET = "validationDatasets/RNAbp/RNABPdatasetFinal.fasta"
#TARGET = "disordered.fasta"
#TARGET = "validationDatasets/prionLikeFromCellPaper/prionLike.fasta"
#TARGET = "validationDatasets/prions/prionDataset.fasta"
TARGET = "validationDatasets/background/backgroundDataset.fasta"
#TARGET= "validationDatasets/disordered_proteins.fasta" #
#TARGET='newPSP.fasta'
ERASE = False
MIN_SEQ_SIZE=150
MAX_SEQ_SIZE=3000

def test(TARGET):
    model=phase_hmm()
    model.fit()


    seqs = leggifasta(PSP)
    targetSeqs = leggifasta(TARGET)
    print ("Found %d proteins in %s" % (len(seqs), PSP))
    print ("Reading PSP...")
    if FAST and os.path.exists("marshalled/phaseSepVects.cPickle") and not ERASE:
        vets = load_pickle("marshalled/phaseSepVects.cPickle")
    else:
        s={}
        for i in seqs.keys()[:]:
            if len(seqs[i])>MAX_SEQ_SIZE :
                continue
            s[i]=seqs[i]
        seqs=s
        vets = model.build_vector(seqs)
        pickle.dump(vets, open("marshalled/phaseSepVects.cPickle","w"))
    phaseSepPreds = model.predict_proba(vets).values()
    ####################################################################

    if  "background" in TARGET:
        targets = targetSeqs.keys()[15000:]

        random.shuffle(targets)
        targets = targets[:2000]

    else:
        targets = targetSeqs.keys()[:]

    print ("Found %d proteins in %s" % (len(targets), TARGET))

    if FAST and os.path.exists("marshalled/"+TARGET.split("/")[-1]+".cPickle") and not ERASE:
        vets = load_pickle("marshalled/"+TARGET.split("/")[-1]+".cPickle")


    else:
        s={}
        for i in targets:
            if len(targetSeqs[i])>MAX_SEQ_SIZE or len(targetSeqs[i])<MIN_SEQ_SIZE:
                continue
            s[i]=targetSeqs[i]
        targets=s

        targets = removeKnownPSP(targets, seqs.keys())  #this removes known PSP

        vets = model.build_vector(targets)
        pickle.dump(vets, open("marshalled/"+TARGET.split("/")[-1]+".cPickle","w"))

    print ('@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@')
    preds = model.predict_proba(vets)
    prediction_diz=preds
    for i in sorted(preds.items(), key=lambda x:x[1], reverse=True)[0:50]:
        print (i[0],i[1],'\\\\')

    targetPreds = preds.values()
    print (targetPreds)
    print ("Using %d proteins" % len(vets))

    ####random.shuffle(preds)
    labels = [1]*len(phaseSepPreds)+[0]*len(targetPreds)
    print ('auc',roc_auc_score(labels, preds))
    print ('pvalue',ranksums(phaseSepPreds, targetPreds))
    print (np.median(phaseSepPreds))

    # plt.figure(figsize=(5,4))
    # plt.title("Ranksums p-value = %.3e\nAUC = %.3f" % (ranksums(phaseSepPreds,targetPreds)[1], roc_auc_score(labels, preds)))
    # plt.grid()
    # plt.ylabel("HMM scores")
    # #plt.violinplot([a,b])
    # plt.boxplot([phaseSepPreds, targetPreds], notch=True)

    # plt.xticks([1,2],["Phase Separation", "IDPs"])
    # plt.tight_layout()
    #plt.savefig(TARGET+".png", dpi=400)
    #plt.show()

    print ('ok')

    return prediction_diz

def removeKnownPSP(t, s):
    #print type(t), type(s)
    r = {}
    c = 0
    for i in t.items():
        if i[0] in s:
            c += 1
            continue
        r[i[0]] = i[1]
    print ("Removed %d prots from targets" % c)
    return r

if __name__ == '__main__':
    test(TARGET)
