# 
# calculate duple orientations by RDC periodicity

import os
import numpy
from math import radians, degrees,sin, cos, pi, sqrt 
from lib.levmar import ddif
from random import random 
from libio import MIO
from copy import copy 

def rdcwave(parameters, dummy, ndata):
    """ Da, R, Theta, Phi, Rho are unknow """
    Da,  R, the, phi, rho_0, delta = parameters
    pm = the, phi, rho_0, delta, Da, R, ndata
    return pm2rdcs_wave(pm)

 
def pm2rdcs_wave(parameters):
    the, phi, rho_0, delta, Da,R, ndata = parameters
    sin_del = sin(delta)
    cos_del = cos(delta)
    sin_2del = 2*sin_del*cos_del
    Ki = (3.0 /16.) * (4.0 + 6.0*R*cos(2.*phi) + R*cos(2.*(phi-the)) - 4. *cos(2.*the) + R*cos(2.*(phi+the)) )*sin_del*sin_del 
    Kii = - Ki
    Kiii = -3.0 * R * cos(the)*sin(2. * phi)* sin_del * sin_del
    Kiv =  0.75  * (-2. + R*cos(2.* phi) )*sin(2.*the)* sin_2del
    Kv = - 3.0 * R* sin(phi)*sin(the)*cos(phi)* sin_2del
    Kvi = (1.0/32.) * ( 4. + 6.* R* cos(2.*phi) - 3. *R*cos(2.*(phi-the)) + 12.*cos(2.*the) - 3.*R*cos(2.*(phi+the)) )*(1. + 3.*cos(2.*delta))
    rdcs = []
    cst =  Kvi + 0.5*(Ki + Kii) 
    for n in ndata:#[1:]
        alf = ALF_PM * (n-1)
        rho = rho_0 + alf 
        rdc  = Da*( 0.5 * Kiii * sin(2.0*rho) + 0.5* (Ki - Kii)*cos(2.*rho) + Kv * sin(rho) + Kiv*cos(rho) + cst )
        rdcs.append(rdc)
    return rdcs 

def SIMRDCWave(parameters, dummy, ndata):
    num_dps = int((len(parameters) -2 )/ 4 ) #
    Da = parameters[-2]
    R = parameters[-1]
    cnt = 0
    nd_id = 0
    rdcs = []
    #print "# all_ndata: ", ndata
    for i in range(num_dps):
        the = parameters[cnt]
        phi = parameters[cnt+1]
        rho = parameters[cnt+2]
        delta = parameters[cnt+3]
        cnt += 4
        num_nd = ndata[nd_id]
        #print "# num_nd: ", num_nd
        nd_id += 1
        nd_end = nd_id + num_nd
        nd = ndata[nd_id:nd_end]
        nd_id = nd_end
        pm = the, phi, rho, delta, Da, R, nd 
        rdcs += pm2rdcs_wave(pm)
        #print "# pm: ", pm 
    #print rdcs 
    return rdcs 


class RDCFit(MIO):
    """ fit rdc by rdc periodicity """
    fit_full_keys  = ["Da", "R", "phi", "the", "rho","delta"]
    
 
    def chkVeglia(self, parameters, vb=0):
        rdcs = rdcwave(parameters,None,self.ndata)
        assert len(self.ydata) == len(rdcs)
        num_data = len(self.ydata)
        sum_sq = 0.
        lns = []
        rdc_difs = []
        for i in range(num_data):
            diff =  self.ydata[i]-rdcs[i]
            rdc_difs.append(diff)
            sum_sq += diff * diff
            lns += [ "# bpid= %3d resid= %3d  exp_rdc= %8.3f  cal_rdc= %8.3f  df= %8.3f " % (self.bp_ids[i], self.dp_rdc_resids[i], self.ydata[i], rdcs[i], self.ydata[i]-rdcs[i])]
        rmsd = sqrt(sum_sq/float(num_data))
        lns += [ "# rmsd= %8.3f,  Da= %8.5f  R= %8.4f   the= %6.1f  phi= %6.1f rho= %6.1f delta= %6.1f  " % \
                 (rmsd, parameters[0], parameters[1], degrees(parameters[2]), degrees(parameters[3]), degrees(parameters[4]), degrees(parameters[5]))]
        if vb: self.prnLns(lns)
        return rmsd, sum_sq, rdc_difs, lns  
    
    
    def fitRDC(self):
        """
        parameters = [ Da, R, the, phi, rho_0]
        """
        lb, ub = self.getLU(self.fitpm)
        all_results = []
        rmsds = []
        for i in range(self.try_num):
            x0 = self.bc2IFitVar(lb,ub)
            results = ddif(func, x0, self.ydata, self.max_iter, data=self.ndata)
            rmsd, sum_sq, rdc_difs, lns = self.chkVeglia(results[0], vb=0)
            #self.prnFitResults(results, rmsd, rho_chk=0)
            all_results.append((results,rdc_difs, rmsd, lns))
            rmsds.append(rmsd)
        return self.sortByRmsd(rmsds, all_results)
    
    def prnFitPM(self):
        """ print out RDC periodicity fitting parameters """
        lns = ["# RDC periodicity fitting parameters "]
        fit_pm = self.fitpm
        pm_keys = fit_pm.keys()
        for _key in self.fit_full_keys:
            if _key in pm_keys:
                if _key in ['Da', "R"]:
                    lns += ["# FitPM:  %-5s= (%7.2f -- %7.2f) " % (_key, fit_pm[_key][0], fit_pm[_key][1]) ]
                else:
                    lns += ["# FitPM:  %-5s= (%7.1f -- %7.1f) " % (_key, degrees(fit_pm[_key][0]), degrees(fit_pm[_key][1])) ]
        lns += ["# fitting try number= %3d  maximum_iteration= %3d " % (self.try_num, self.max_iter) ]
      