# python modules
import Bio.PDB
import numpy
import os
import gzip
import re
import collections as col
import copy
import argparse
import tarfile
import ast

import os
import sys
import modeller
import re
import pylab

from modeller import *
from modeller.scripts import complete_pdb
from modeller.optimizers import conjugate_gradients, molecular_dynamics, actions

# our modules
from functions import *
from classes import *
from reduce_inputs_func import *
from DOPE_profile import *
import utilities
import re

def refine_model(outputs, log_files, stats, refined_models, dope_profile):
    """Refines a model previously generated with the program according to the optimization parameters defined by
     Modeller."""
    blockPrint()
    import shutil
    env = environ()
    env.io.atom_files_directory = ['../atom_files']
    env.edat.dynamic_sphere = True

    env.libs.topology.read(file='$(LIB)/top_heav.lib')
    env.libs.parameters.read(file='$(LIB)/par.lib')

    dl = os.listdir(outputs)

    files_list = []
    dict_finalen = {}
    list_finalen = []

    for file in dl:
        if file.endswith(".pdb"):
            files_list.append(file)

    enablePrint()
    print("Analyzing ", len(files_list), "model(s)...")
    blockPrint()

    for file in files_list:
        blockPrint(file=log_files + str(file[:-4]) + '.log')
        print("\n", "##########", file, "##########", "\n")
        code = outputs + str(file)
        mdl = complete_pdb(env, code)
        atmsel = selection(mdl)
        cg = conjugate_gradients(output='REPORT')

        trcfil = open(stats + str(file) + '.stats', 'w')

        # Run CG on the all-atom selection; write stats every 5 steps
        cg.optimize(atmsel, max_iterations=20, actions=actions.trace(10, trcfil))
        mpdf2 = atmsel.energy()
        mdl.write(file=refined_models + str(file) + '.B')
        print("Final energy: ", mpdf2[0])
        dict_finalen[file] = mpdf2[0]
        list_finalen.append(mpdf2[0])

    for element in os.listdir(refined_models):
        if element.endswith('.B'):
            shutil.copy2(refined_models+element, outputs)

    enablePrint()

    print("The minimum final energy for all the models is:", min(list_finalen), "and corresponds to:",
          list(dict_finalen.keys())[list(dict_finalen.values()).index(min(list_finalen))] + '.B')

    blockPrint()
    return "Refined"

def refined_vs_notrefined(models_dir, dope_profile):
    """Creates a comparison energy plot between the model generated by the program and the refined one."""
    env = environ()
    env.io.atom_files_directory = [models_dir]

    mdl_list = []
    aln = modeller.alignment(env)
    code_list = []

    mdl_nr_list = []
    aln_nr = modeller.alignment(env)
    code_list_nr = []

    for file in os.listdir(models_dir):
        name = file
        if name.endswith('.B'):
            mdl = modeller.model(env)
            mdl.read(file = file)
            code = str(file)
            code_list.append(code)
            s = selection(mdl)
            s.assess_dope(output='ENERGY_PROFILE NO_REPORT', file= models_dir + code + '.profile',
                  normalize_profile=True, smoothing_window=15)
            mdl_list.append(mdl)
            aln.append_model(mdl, align_codes = code, atom_files = code)
            aln.write(file=dope_profile+'build_profile_ref.ali', alignment_format='PIR')
        else:
            mdl_nr = modeller.model(env)
            mdl_nr.read(file = file)
            code = str(file)
            print (code)
            code_list_nr.append(code)
            t = selection(mdl_nr)
            t.assess_dope(output='ENERGY_PROFILE NO_REPORT', file= models_dir + code + '.profile',
                  normalize_profile=True, smoothing_window=15)
            mdl_nr_list.append(mdl_nr)
            aln_nr.append_model(mdl_nr, align_codes = code, atom_files = code)
            aln_nr.write(file=dope_profile+'build_profile_notref.ali', alignment_format='PIR')

    if len(mdl_nr_list) == len(mdl_list):
        for a, b, c, d in zip(mdl_nr_list, mdl_list, code_list_nr, code_list):
            model1 = get_profile(models_dir + c + ".profile", aln_nr[str(c)])
            model2 = get_profile(models_dir + d + ".profile", aln[str(d)])
            pylab.figure(1, figsize=(30,18))
            pylab.xlabel('Alignment position', fontsize = 20)
            pylab.ylabel('DOPE per-residue score', fontsize = 20)
            pylab.plot(model1, color='red', linewidth=2, label='Model')
            pylab.plot(model2, color='green', linewidth=2, label='Optimized model')
            pylab.legend(fontsize = 20)
            pylab.savefig(dope_profile + c + '.dope_profile.jpg', dpi=100)
            pylab.close()

    return ("Comparison energy plot between refined and not refined model has been created here:\n%s\n" % (dope_profile))