# This file is part of the xlms-tools package
#
# Copyright (c) 2023 - Topf Lab, Leibniz-Institut für Virologie
# and Center for Data and Computing in Natural Sciences, Universität Hamburg
# Hamburg, Germany.
#
# This module was developed by:
#   Karen Manalastas-Cantos    <karen.manalastas-cantos AT cssb-hamburg.de>

from Bio.PDB import *
from xlms_tools.depth import *
from xlms_tools.parsers import *
from xlms_tools.visualization import *
import numpy as np
import time

class StructModel:
    def __init__(self, infile, depths=None, verbose=True):
        if verbose:
            print ("MODEL:", infile)

        # if no depths computed, we are starting from a pdb/cif file
        if depths == None:
            self.filename = infile
            self.biopdbstruct = parsebiopystructure(infile)
            self.residuedepths, self.residues, self.edt, self.nearestwaters, self.voxelized, self.atomlist, self.translatedatomcoords = computedepth(self.biopdbstruct, return_indices=True, verbose=verbose)

        # depth has been precomputed, probably api access
        else:
            self.filename = None
            self.biopdbstruct = infile
            self.residuedepths = depths
            self.residues, self.edt, self.nearestwaters = None, None, None
            self.voxelized, self.atomlist, self.translatedatomcoords = None, None, None
        
        self.maxCAdist = maxCAdistfromstructure(self.biopdbstruct)
        self.maxdepth = max(self.residuedepths.values())
        self.xlscores = {}
        self.avexlscore = 0
        self.avemlscore = 0
        self.depviols = 0
        self.distviols = 0
    
    def inmodel(self, link):
        found = False
        if link.__class__.__name__ == 'Monolink':
            if self.biopdbstruct[0][link.chainid][link.rnum].__contains__('CA'):
                return True
        else:   #if crosslink
            if self.biopdbstruct[0][link.siteA.chainid][link.siteA.rnum].__contains__('CA'):
                if self.biopdbstruct[0][link.siteB.chainid][link.siteB.rnum].__contains__('CA'):
                    return True
        return found
            
    def scoremodel(self, xlmslist, verbose=True):
        start = time.time()
        
        # accumulators for denominators
        totalmlweight = 0
        totalxlweight = 0
        xlweightbychain = {}
        
        # score each link
        for link in xlmslist:
            if self.inmodel(link):
                if link.__class__.__name__ == 'Monolink':
                    mlscore, depth, viol = link.score(self)
                    self.avemlscore += mlscore
                    totalmlweight += link.weight
                    if viol:
                        self.depviols += 1
                else:   # if crosslink
                    xlscore, dist, viol = link.score(self)
                    self.avexlscore += xlscore
                    totalxlweight += link.weight
                    chid = link.chains() 
                    if chid not in self.xlscores:
                        self.xlscores[chid] = xlscore
                        xlweightbychain[chid] = link.weight
                    else:
                        self.xlscores[chid] += xlscore
                        xlweightbychain[chid] += link.weight
                    if viol:
                        self.distviols += 1

        # get score averages
        if totalmlweight != 0:
            self.avemlscore /= totalmlweight
        if totalxlweight != 0:
            self.avexlscore /= totalxlweight
        for key in self.xlscores:
            self.xlscores[key] /= xlweightbychain[key]
        if verbose:
            print (f"-- Scoring model: {time.time()-start:.3f}s\n")    

    def printscores(self, ptr):
        
        # order output file columns
        intra = [i for i in self.xlscores if len(i) == 1]
        inter = [i for i in self.xlscores if len(i) == 2]
        intra.sort()
        inter.sort()
    
        
        ptr.write(f'{os.path.basename(self.filename)[:-4]}\t')
        if self.avemlscore != 0:
            ptr.write(f'{self.avemlscore:.3f}\t')
        else:
            ptr.write('--\t')
            
        
        for key in intra:
        	ptr.write(f'{self.xlscores[key]:.3f}\t')
        for key in inter:
        	ptr.write(f'{self.xlscores[key]:.3f}\t')
    
        ptr.write(f'{self.depviols}\t')
        ptr.write(f'{self.distviols}\t')
        ptr.write(f'{self.avexlscore:.3f}')            
        ptr.write('\n')
    
    def printchimeraxfiles(self, cxdir, cxcptr, xlmslist, modelid):
        # open pseudobond files
        g, m, b = createpb(self.filename, cxdir)
        
        # go through crosslinks and monolinks
        for link in xlmslist:
            if self.inmodel(link):
                if link.__class__.__name__ == 'Monolink':
                    # add visualization of monolink to ChimeraX command file
                    cxcptr.write(f'show #{modelid}/{link.chainid}:{link.rnum} atoms\n')
                    cxcptr.write(f'color #{modelid}/{link.chainid}:{link.rnum} {link.color(self)} atoms\n')
                    
                else:   # if crosslink
                    dist = link.eudist(self)
                    pos1 = f'{link.siteA.chainid}:{link.siteA.rnum}'
                    pos2 = f'{link.siteB.chainid}:{link.siteB.rnum}'
                    
                    if dist <= link.fiftypercent():
                        g.write(f'#{modelid}/{pos1}@ca #{modelid}/{pos2}@ca\n')
                    elif dist <= link.cutoff():
                        m.write(f'#{modelid}/{pos1}@ca #{modelid}/{pos2}@ca\n')
                    else:
                        b.write(f'#{modelid}/{pos1}@ca #{modelid}/{pos2}@ca\n')

        # close pseudobond files
        g.close()
        m.close()
        b.close()
                    
    def outputdepth(self):
        with open(self.biopdbstruct.id+'-residue.depth', 'w') as f:
            for key in self.residuedepths.keys():
                f.write(f'{key}\t{self.residues[key]}\t{self.residuedepths[key]}\n')    
    
class Crosslink:
    def __init__(self, monolinkA, monolinkB, weight=1.0):
        self.siteA = monolinkA
        self.siteB = monolinkB
        self.weight = float(weight)

    def linkdetails(self):
        self.siteA.linkdetails()
        self.siteB.linkdetails()
        print (self.weight)
    
    def chains(self):
        if self.siteA.chainid == self.siteB.chainid:
            return self.siteA.chainid
        else:
            return self.siteA.chainid+self.siteB.chainid

    def probability(self, dist):
        if self.siteA.linker == 'BS3/DSS':    # default situation, BS3 or DSS
            return 1/(1+np.exp(0.33*dist-7))    # this is the best from the benchmark
        elif self.siteA.linker == 'DSBSO':
            return logistic_function(dist, 1.07, 0.18, 20.25)
        elif self.siteA.linker == 'PhoX':
            return logistic_function(dist, 1.05, 0.34, 15.50)
        else:
            print ('No information about linker. Using distance distribution for BS3/DSS')
            return 1/(1+np.exp(0.33*dist-7))

    def cutoff(self):
        cutoffs = {'BS3/DSS':33, 'DSBSO':40, 'PhoX':30}
        if self.siteA.linker in cutoffs:
            return cutoffs[self.siteA.linker]
        else:
            print ('No information about linker. Using distance cutoff for BS3/DSS')
            return cutoffs['BS3/DSS']
                
    def eudist(self, model):
        ca1 = model.biopdbstruct[0][self.siteA.chainid][self.siteA.rnum]['CA']
        ca2 = model.biopdbstruct[0][self.siteB.chainid][self.siteB.rnum]['CA']
        return ca1-ca2
    
    def fiftypercent(self):
        cutoffs = {'BS3/DSS':21, 'DSBSO':21, 'PhoX':16}
        if self.siteA.linker in cutoffs:        
            return cutoffs[self.siteA.linker]  
        else:
            print ('No information about linker. Using distance cutoff for BS3/DSS')
            return cutoffs['BS3/DSS']
         
    def score(self, model):
        viol = False
        dep1 = self.siteA.depth(model)
        dep2 = self.siteB.depth(model)
        mp1 = self.siteA.probability(dep1)
        mp2 = self.siteB.probability(dep2)
        dist = self.eudist(model)
        distp = self.probability(dist)
        score = mp1*mp2*distp
        
        if dep1 >= self.siteA.cutoff():
            score = -dep1/(4*model.maxdepth)
        if dep2 >= self.siteB.cutoff():
            score = -dep2/(4*model.maxdepth)
        if dist > self.cutoff():
            score = -dist/(2*model.maxCAdist)
            viol = True
        score *= self.weight
        return score, dist, viol
        
class Monolink:
    def __init__(self, chainid, rnum, linker='BS3/DSS', weight=1.0):
        self.chainid = chainid
        self.rnum = int(rnum)
        self.linker = linker
        self.weight = float(weight)
    
    def linkdetails(self):
        print (self.chainid, self.rnum, self.linker, self.weight)
            
    def probability(self, depth):
        if self.linker == 'BS3/DSS':
            tagcoefs = np.array([1.20875756e+03, 5.31961163e+00, 6.43465155e-04])
            lyscoefs = np.array([4.44470936e+02, 4.65135510e+00, 1.69863402e-03])
            if func(depth, *lyscoefs) > 0:
                return func(depth, *tagcoefs)/func(depth, *lyscoefs)
            else:
                return 0
        elif self.linker == 'DSBSO':
            tagcoefs = np.array([2583, 5.38])
            lyscoefs = np.array([318, 3.94])
            if func2(depth, *lyscoefs) > 0:
                return func2(depth, *tagcoefs)/func2(depth, *lyscoefs)
            else:
                return 0
        elif self.linker == 'PhoX':
            tagcoefs = np.array([2896, 5.47])
            lyscoefs = np.array([288, 3.88])
            if func2(depth, *lyscoefs) > 0:
                return func2(depth, *tagcoefs)/func2(depth, *lyscoefs)
            else:
                return 0

    def cutoff(self):
        cutoffs = {'BS3/DSS':15, 'DSBSO':15, 'PhoX':17}
        if self.linker in cutoffs:        
            return cutoffs[self.linker]  
        else:
            print ('No information about linker. Using depth cutoff for BS3/DSS')
            return cutoffs['BS3/DSS']
        
    def depth(self, model):
        return model.residuedepths[f'{self.chainid}:{self.rnum}']
    
    def color(self, model):
        depth = self.depth(model)
        color = 'blue'
        
        if self.linker == 'BS3/DSS':
            if depth > 6.25:
                color = 'red'
        elif self.linker == 'DSBSO':
            if depth > 6.25:
                color = 'red'
        elif self.linker == 'PhoX':
            if depth > 6.25:
                color = 'red'
        else:
            if depth > 6.25:
                color = 'red'            
        return color
        
    def score(self, model):
        viol = False
        depth = self.depth(model) 
        if depth >= self.cutoff():
            score = -(depth/model.maxdepth)
            viol = True
        else:
            score = self.probability(depth)
        score *= self.weight
        return score, depth, viol
        
def getxlmslist(xlmslist, linker='BS3/DSS'):
    xllist = []
    unique = []
    with open(xlmslist, 'r') as f:
        for ln in f:
            link, linktuple = parsexlmsline(ln.strip(), linker=linker)
            if linktuple not in unique:
                unique.append(linktuple)
                xllist.append(link)
    return xllist

def parsexlmsline(line, linker='BS3/DSS', weight=1.0):
    tmp = line.split()
    if len(tmp) > 1:    #if there's a score
        weight = float(tmp[1])
        buf = tmp[0].split('|')
    else:
        buf = line.split('|')
        
    if len(buf) == 2:   # if monolink
        res1 = int(buf[0])
        ch1 = buf[1]
        site1 = Monolink(ch1, res1, linker=linker, weight=weight)
        return site1, (ch1, res1)
        
    elif len(buf) == 4:   # if crosslink
        res1 = int(buf[0])
        ch1 = buf[1]
        site1 = Monolink(ch1, res1, linker=linker)
        res2 = int(buf[2])
        ch2 = buf[3]
        site2 = Monolink(ch2, res2, linker=linker)
        if (ch1, res1) <= (ch2, res2):
            return Crosslink(site1, site2, weight=weight), (ch1, res1, ch2, res2)
        else:
            return Crosslink(site2, site1, weight=weight), (ch2, res2, ch1, res1)

def linkmetrics(structure, linktuple, depths, linkweight=1.0, linker='BS3/DSS'):
    # initialize structure model
    model = StructModel(structure, depths=depths)
    
    # initialize crosslink or monolink
    if len(linktuple) == 2:  #if a monolink
        link = Monolink(linktuple[1], linktuple[0], linker=linker, weight=linkweight)
    else:   #if a crosslink
        ml1 = Monolink(linktuple[1], linktuple[0], linker=linker)
        ml2 = Monolink(linktuple[3], linktuple[2], linker=linker)
        link = Crosslink(ml1, ml2, weight=linkweight)
    
    # compute score and depth/distance
    if model.inmodel(link): # link found in model
        score, d, viol = link.score(model)
        return score, d
    else:   # link not found
        print ("Link coordinates not found in model")
        return None, None
            
def func(x, a, b, c):
    return (a * (x**-b)) + c

def func2(x, a, b):
    return (a * (x**-b))

def logistic_function(x, L, k, x0):
    return L / (1 + np.exp(k * (x - x0)))

