#!/bin/python
#
# Create a plot of free energy vs temperature for a polymorph
#
# Copyright Michael R. Shirts, University of Virginia, 2014
#
from __future__ import print_function
import numpy as np
import pymbar # multistate Bennett acceptance ratio
from pymbar import timeseries # timeseries analysis
from optparse import OptionParser # for parsing command-line options
import mdtraj as md
import MBARBootstrap # Bootstrapping algorithm
import random
import os
import usefulFuncs #Useful math functions
import Harvist #Hamiltonian Reweighting Visualization Toolkik
import pdb
import sys
import panedr
import matplotlib
import subprocess
import matplotlib.pyplot as plt

font = {'family': 'normal',
        'weight': 'normal',
        'size': 16}

def old_systems_dictionary(potential, molecule):
    Polys = dict()  # Name of the polymorphs
    refTs = dict()  # Reference temperatures for the PSCP for each system
    refdGs = dict()  # Reference free energies for the PSCP for each system
    refddGs = dict()  # Reference uncertainties for the PSCP for each system
    refdUs = dict()  # Reference lattice minima for each system
    absolutedUs = dict()  # Absolute lattice energy for each system

    # Oplsaa
    if potential == "oplsaa":
        Potentials = ['oplsaa']
        PotNAME = 'OPLS'
        Charges = ['0.1150', '0.1150', '0.0700', '0.0850', '0.1000', '0.1150', '0.1300', '0.1450', '0.1600', '0.1750',
                   '0.1900', '0.1000', '0.0900', '0.0800', '0.0700', '0.0600', '0.0500', '0.0400', '0.0300', '0.0200',
                   '0.0100']
        Chargenames = ['', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '',
                       '', '', '', '', '', '', '', '']
        #Chargenames=['C01150', 'C01150', 'C00700', 'C00850', 'C01000', 'C01150', 'C01300', 'C01450', 'C01600', 'C01750', 'C01900', 'C00900', 'C00800', 'C00700', 'C00600', 'C00500', 'C00400', 'C00300', 'C00200', 'C00100']
        PotNAMES=['OPLS']
        SimNAMES=['GRO']
        Temperatures = np.array([50, 100, 200, 300])

        #Temperatures=np.array([45,50]) #Overlap Check
        #Temperatures=np.array([10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,70,80,90,100,110,120,130,140])#,80,90,100,110])#Zzzvye
        #Temperatures=np.array([10,15,20,25,30,35,40,45,50,60,70,80,90,100,110])#Zzzvye
        #Temperatures = np.array([10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280,290,300])#,320,340])#,360,380,400,420,440,460,480,500])#,320,340,360,380,400]) #benzene
        #Temperatures=np.array([10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280,290,300])#,320,340])#,360,380,400,420,440,460,480,500])#,320,340,360,380,400]) #kobfud
        #Temperatures=np.array([10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280,290,300])#,320,340,360,380,400,420,440,460,480,500])#,320,340,360,380,400]) #cbmzpn
        #Temperatures=np.array([10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280,290,300])#,320,340])#,360,380,400,420,440,460,480,500])#,320,340,360,380,400]) #melfit
        #Temperatures=np.array([10,20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280,290,300])#,320,340,360,380,400,420,440,460,480,500])
        #Temperatures=np.array([100,130,140,160,180,200,220,240,260,280,300,340,360,380,400])
        #Temperatures=np.array([100,120,140,160,180,200]) #MelfitTest
        #ExtraPressures=np.array([5000,15000,35000]) #benzene
        #Temperatures=np.array([40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280,290,300]) #imazol
        #ExtraPressures=np.array([5000,10000,16000,20000,26000,31000,36000,40000]) #imazol
        #Temperatures=np.array([20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,220,240,260,280,300,300,300,300,300,300,300,300]) #acetac
        #ExtraPressures=np.array([5000,5000,5000,5000,5000,10000,16000,20000,26000,31000,36000,40000]) #acetac
        #Temperatures=np.array([20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250,200,200,200,200,220,240,260,280,300,300,300,300,300,300,300,300]) #formam
        #ExtraPressures=np.array([1500,2000,5000,10000,10000,10000,10000,10000,5000,10000,15000,20000,26000,31000,36000,40000]) #formam
        #Temperatures=np.array([20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,220,240,250,270,290,300,320,340,360,380,400,420,440,460,480,500]) #glycin
        ExtraPressures = []
        Pressures = np.ones(len(Temperatures), int)
        Pressures[len(Pressures) - len(ExtraPressures): len(Pressures)] = ExtraPressures
        #Temperatures=np.array([200,200,200,200,200])
        #Pressures=np.array([1,500,1000,1500,2000,2500,3000,3500,4000,4500,5000])
        refPot = 0

        #Benzene
        Polys['benzene']=['p1', 'p2', 'p3']
        refTs['benzene']=200
        refdGs['benzene']=[0.000, 0.185, 0.306]
        refddGs['benzene']=[0.000, 0.019, 0.019]
        refdUs['benzene'] = [0.000, 0.267, 0.240]
        absolutedUs['benzene'] = [-5.624, -5.362, -5.380]

        #Formam
        Polys['formam']=['p1', 'p2']
        refTs['formam']=200
        refdGs['formam']=[0.000, 0.118]
        refddGs['formam']=[0.005, 0.005]
        refdUs['formam'] = [0.000, 0.343]
        absolutedUs['formam'] = [-41.04733, -40.7076]

        #Acetac
        Polys['acetac']=['p1', 'p2']
        refTs['acetac']=200
        refdGs['acetac']=[0.000, 0.001]
        refddGs['acetac']=[0.005, 0.005]
        refdUs['acetac'] = [0.000, 0.020]
        absolutedUs['acetac'] = [-36.4858, -36.466]

        #Imazol
        Polys['imazol']=['p1', 'p2']
        refTs['imazol']=100
        refdGs['imazol']=[0.000, -0.436]
        refddGs['imazol']=[0.005, 0.005]
        refdUs['imazol'] = [0.000, -0.464]
        absolutedUs['imazol'] = [-15.399, -15.863]

        #Glycin
        Polys['glycin']=['p1', 'p2', 'p4']#, 'p3']
        refTs['glycin']=200
        refdGs['glycin']=[0.000, -0.034, -0.067]#, 1.087]
        refddGs['glycin']=[0.005, 0.005, 0.005]#, 0.005]
        refdUs['glycin'] = [0.000, 0.117, 0.288]#, 1.235]
        absolutedUs['glycin'] = [-131.950, -131.778, -131.609]#, -130.715]

        #Bismev
        Polys['bismev']=['p3', 'p1', 'p2']
        refTs['bismev']=200
        refdGs['bismev']=[0.000, -1.474, -0.186]
        refddGs['bismev']=[0.005, 0.005, 0.005]#, 0.005]
        refdUs['bismev'] = [0.000, -1.459, -0.245]#, 1.235]
        absolutedUs['bismev'] = [-64.3467, -65.8055, -64.5916]#, -130.715]

        #Cbmzpn
        Polys['cbmzpn']=['p3', 'p1']
        refTs['cbmzpn']=200
        refdGs['cbmzpn']=[0.000, 0.195]
        refddGs['cbmzpn']=[0.0, 0.005]
        refdUs['cbmzpn'] = [0.000, 0.565]
        absolutedUs['cbmzpn'] = [-34.3933, -33.8283]

        #Hxacan
        Polys['hxacan']=['p2', 'p1']
        refTs['hxacan']=200
        refdGs['hxacan']=[0.000, -0.619]
        refddGs['hxacan']=[0.005, 0.005]
        refdUs['hxacan'] = [0.000, -0.446]
        absolutedUs['hxacan'] = [-46.965, -47.411]

        #Kobfud
        Polys['kobfud']=['p1', 'p2']
        refTs['kobfud']=200
        refdGs['kobfud']=[0.000, 1.832]
        refddGs['kobfud']=[0.005, 0.005]
        refdUs['kobfud'] = [0.000, 2.255]
        absolutedUs['kobfud'] = [-49.974, -47.719]

        ##Pyrzin
        #Polys['pyrzin']=['p2']#, 'p3', 'p1', 'p5', 'p2']
        #refTs['pyrzin']=200
        #refdGs['pyrzin']=[0.000] #, -0.483, -0.689, -1.055, -1.056]
        #refddGs['pyrzin']=[0.005] #, 0.005, 0.005, 0.005, 0.005]
        #refdUs['pyrzin'] = [0.000] #, -0.176, -0.695, -1.246, 0.951]
        #absolutedUs['pyrzin'] = [-13.348] #, -14.475, -14.993, -15.545, -13.348]

        #Pyrzin
        Polys['pyrzin']=['p4', 'p3', 'p1', 'p5', 'p2']
        refTs['pyrzin']=200
        refdGs['pyrzin']=[0.000, -0.483, -0.689, -1.055, -1.056]
        refddGs['pyrzin']=[0.005, 0.005, 0.005, 0.005, 0.005]
        refdUs['pyrzin'] = [0.000, -0.176, -0.695, -1.246, 0.951]
        absolutedUs['pyrzin'] = [-14.299, -14.475, -14.993, -15.545, -13.348]

        #Qopbed
        Polys['qopbed']=['p1', 'p3']#, 'p2']
        refTs['qopbed']=200
        refdGs['qopbed']=[0.000, 0.878] #1.699]
        refddGs['qopbed']=[0.005, 0.005]
        refdUs['qopbed'] = [0.000, 1.235]# 2.778]
        absolutedUs['qopbed'] = [-41.135, -39.900] # -38.357]

        #Zzzvye
        Polys['zzzvye']=['p3', 'p1']
        refTs['zzzvye']=50
        refdGs['zzzvye']=[0.000, 0.190]
        refddGs['zzzvye']=[0.005, 0.005]
        refdUs['zzzvye'] = [0.000, 0.351]
        absolutedUs['zzzvye'] = [11.103, 11.454]

        #Resora
        Polys['resora']=['p1', 'p2']
        refTs['resora']=200
        refdGs['resora']=[0.000, -0.337]
        refddGs['resora']=[0.005, 0.005]
        refdUs['resora'] = [0.000, -0.278]
        absolutedUs['resora'] = [-22.506, -22.784]

        #Zzzpro
        Polys['zzzpro']=['p1', 'p3']#, 'p2']
        refTs['zzzpro']=200
        refdGs['zzzpro']=[0.000, 0.223]#, 0.409]
        refddGs['zzzpro']=[0.01, 0.01]#, 0.01]
        refdUs['zzzpro'] = [0.000, 0.258]#, 0.480]
        absolutedUs['zzzpro'] = [-19.393, -19.135]#, -18.913]

        #Zzzpus
        Polys['zzzpus']=['p6', 'p2'] #, 'p1']#, 'p3', 'p7']
        refTs['zzzpus']=200
        refdGs['zzzpus']=[0.000, -1.713] #, -1.713]#, -1.713, -1.713]
        refddGs['zzzpus']=[0.00, 0.028] #, 0.02]#, 0.02, 0.02]
        refdUs['zzzpus'] = [0.000, -2.226] #, 1.4102]#, -2.249, -3.067]
        absolutedUs['zzzpus'] = [-84.917, -87.143] #, -83.507]#, -87.166, -87.985]

        ##Melfit
        #Polys['melfit']=['p7', 'p6']#, 'p8']
        #refTs['melfit']=200
        #refdGs['melfit']=[0.000, 0.52]#, 0.52]
        #refddGs['melfit']=[0.00, 0.015]#, 0.02]
        #refdUs['melfit'] = [0.000, 0.596]#, 1.187]
        #absolutedUs['melfit'] = [-12.456, -11.818]#, -11.268]

        #Melfit
        Polys['melfit']=['p7', 'p6', 'p8', 'p1', 'p5']
        refTs['melfit']=200
        refdGs['melfit']=[0.000, 0.52, 0.52, 0.52, 0.52]
        refddGs['melfit']=[0.02, 0.02, 0.02, 0.02, 0.02]
        refdUs['melfit'] = [0.000, 0.637, 1.187, 1.966, 0.927]
        absolutedUs['melfit'] = [-12.456, -11.818, -11.268, -10.489, -11.529]

        ##Melfit
        #Polys['melfit']=['p7', 'p6', 'p1', 'p5']#, 'p2', 'p3']#, 'p4']
        #refTs['melfit']=200
        #refdGs['melfit']=[0.000, 0.279, 1.581, 0.542]#, -0.962, 1.361]#, -0.679]
        #refddGs['melfit']=[0.02, 0.02, 0.02, 0.02]#, 0.02, 0.02]#, 0.02]
        #refdUs['melfit'] = [0.000, 0.279, 1.581, 0.542]#, -0.962, 1.361]#, -0.679]
        #absolutedUs['melfit'] = [-12.070, -11.792, -10.489, -11.529]#, -13.032, -10.709]#, -12.749]

        ##Bedmig
        #Polys['bedmig']=['p1', 'p2', 'p6', 'p4', 'p7', 'p3', 'p5']
        #refTs['bedmig']=200
        #refdGs['bedmig']=[0.000, 1.678, -0.372, -2.878, 1.955, -0.322, -1.146]
        #refddGs['bedmig']=[0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02]
        #refdUs['bedmig'] = [0.000, 3.803, -0.649, -1.217, 2.772, 2.777, 5.638]
        #absolutedUs['bedmig'] = [-92.82, -89.01, -93.46, -94.03, -90.03, -90.0378, -87.1763]

        #Bedmig
        Polys['bedmig']=['p1', 'p7']
        refTs['bedmig']=200
        refdGs['bedmig']=[0.000, 1.955]
        refddGs['bedmig']=[0.00, 0.016]
        refdUs['bedmig'] = [0.000, 2.789]
        absolutedUs['bedmig'] = [-92.814, -90.025]

        ##Bedmig
        #Polys['bedmig']=['p5']
        #refTs['bedmig']=200
        #refdGs['bedmig']=[0.000]
        #refddGs['bedmig']=[0.00]
        #refdUs['bedmig'] = [0.000]
        #absolutedUs['bedmig'] = [-87.1763]

        ##Bedmig
        #Polys['bedmig']=['p1', 'p7', 'p3', 'p5', 'p6']
        #refTs['bedmig']=200
        #refdGs['bedmig']=[0.000, 1.955, -0.322, -1.146, -0.372]
        #refddGs['bedmig']=[0.02, 0.02, 0.02, 0.02, 0.02]
        #refdUs['bedmig'] = [0.000, 2.772, 2.777, 5.638, -0.649]
        #absolutedUs['bedmig'] = [-92.82, -90.03, -90.0378, -87.1763, -93.46]

        #Cafine
        Polys['cafine']=['p1', 'p2']
        refTs['cafine']=200
        refdGs['cafine']=[0.000, 1.846]
        refddGs['cafine']=[0.005, 0.005]
        refdUs['cafine'] = [0.000, 2.375]
        absolutedUs['cafine'] = [-28.737, -26.362]

        SystemNAME = "OPLSAA"

    # Amoeba Reweighting
    if potential == "amoeba09":
        #Potentials=['amoeba09','designeda']#, 'designeda', 'designeda', 'designeda', 'designeda', 'designeda']
        Potentials=['amoeba09']
        PotNAME='AMO'
        Charges=['0.1150', '0.1150', '0.0700', '0.0850', '0.1000', '0.1150', '0.1300', '0.1450', '0.1600', '0.1750', '0.1900', '0.1000', '0.0900', '0.0800', '0.0700','0.0600','0.0500','0.0400','0.0300','0.0200','0.0100']
        Chargenames=['', '', 'C00700', 'C00850', 'C01000', 'C01150', 'C01300', 'C01450', 'C01600', 'C01750', 'C01900', 'C00900', 'C00800', 'C00700', 'C00600', 'C00500', 'C00400', 'C00300', 'C00200', 'C00100']
        PotNAMES=['AMO', 'DESA']
        SimNAMES=['TIN', 'GRO']
        #PotNAMES=['AMO']
        #SimNAMES=['TIN']
        refPot=0
        Temperatures=np.array([30,40,50,60,70,80,90,100,110,130,140,150,160,170,180,190,200,210,220,230,240,250])
        Pressures=np.ones(len(Temperatures),int);
        #Temperatures=np.array([60,100,140,200])

        #Benzene
        Polys['benzene']=['p1', 'p2', 'p3']
        refTs['benzene']=200
        refdGs['benzene']=[0.000, -0.362, -0.291]
        #refdGs['benzene']=[0.000, 0.390, 0.330]
        refddGs['benzene']=[0.000, 0.031, 0.032]
        refdUs['benzene'] = [0.000, 0.781, 0.643]
        absolutedUs['benzene'] = [-3.587, -2.807, -2.944]

        #Formam
        Polys['formam']=['p1', 'p2']
        refTs['formam']=100
        refdGs['formam']=[0.000, 0.293]
        refddGs['formam']=[0.000, 0.020]
        refdUs['formam'] = [0.000, 0.487]
        absolutedUs['formam'] = [-26.237, -25.750]

        #Acetac
        Polys['acetac']=['p1', 'p2']
        refTs['acetac']=50
        refdGs['acetac']=[0.000, -0.106]
        refddGs['acetac']=[0.02, 0.02]
        refdUs['acetac'] = [0.000, -0.136]
        absolutedUs['acetac'] = [-34.028, -34.163]

        #Imazol
        Polys['imazol']=['p1', 'p2']
        refTs['imazol']=200
        refdGs['imazol']=[0.000, 0.125]
        refddGs['imazol']=[0.020, 0.020]
        refdUs['imazol'] = [0.000, 0.201]
        absolutedUs['imazol'] = [-23.186, -22.985]

        #Glycin
        Polys['glycin']=['p1', 'p2', 'p3']
        refTs['glycin']=10
        refdGs['glycin']=[0.000, 0.163, 1.222]
        refddGs['glycin']=[0.000, 0.000, 0.000]
        refdUs['glycin'] = [0.000, 0.163, 1.222]
        absolutedUs['glycin'] = [0.000, 0.000, 0.000]

        SystemName = "AMOEBA"

    #Gromos
    if potential == "gromos":
        Potentials=['gromos54a7']
        PotNAME='GROM'
        Charges=['0.1150', '0.1150', '0.0700', '0.0850', '0.1000', '0.1150', '0.1300', '0.1450', '0.1600', '0.1750', '0.1900', '0.1000', '0.0900', '0.0800', '0.0700','0.0600','0.0500','0.0400','0.0300','0.0200','0.0100']
        Chargenames=['C01150', 'C01150', 'C00700', 'C00850', 'C01000', 'C01150', 'C01300', 'C01450', 'C01600', 'C01750', 'C01900', 'C00900', 'C00800', 'C00700', 'C00600', 'C00500', 'C00400', 'C00300', 'C00200', 'C00100']
        PotNAMES=['GROM']
        SimNAMES=['GRO']
        Temperatures=np.array([10,20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200,210,220,230,240,250])
        refPot=0

        #Benzene
        Polys['benzene']=['p1', 'p2', 'p3']
        refTs['benzene']=200
        refdGs['benzene']=[0.000, 0.251, 0.345]
        refddGs['benzene']=[0.000, 0.019, 0.018]
        refdUs['benzene'] = [0.000, 0.314, 0.328]
        absolutedUs['benzene'] = [-10.184, -9.869, -9.853]

        SystemName = "GROMOS"

    Polymorphs = Polys[molecule]
    refT = refTs[molecule]
    refdG = refdGs[molecule]
    refddG = refddGs[molecule]
    refdU = refdUs[molecule]
    absolutedU = absolutedUs[molecule]
    return Polymorphs, refT, refdG, refddG, refdU, absolutedU

def get_potential_info(potential):
    if potential == 'oplsaa':
        SimNAMES=['GRO']
        Chargenames = ['', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '',
                       '', '', '', '', '', '', '', '', '']
        PotNAMES = ['OPLS']
    elif potential == 'amoeba09':
        SimNAMES=['TIN', 'GRO']
        Chargenames = ['', '', 'C00700', 'C00850', 'C01000', 'C01150', 'C01300', 'C01450', 'C01600', 'C01750', 'C01900',
                       'C00900', 'C00800', 'C00700', 'C00600', 'C00500', 'C00400', 'C00300', 'C00200', 'C00100']
        PotNAMES = ['AMO', 'DESA']
    elif potential == "gromos":
        SimNAMES = ['GRO']
        Chargenames = ['C01150', 'C01150', 'C00700', 'C00850', 'C01000', 'C01150', 'C01300', 'C01450', 'C01600',
                       'C01750', 'C01900', 'C00900', 'C00800', 'C00700', 'C00600', 'C00500', 'C00400', 'C00300',
                       'C00200', 'C00100']
        PotNAMES = ['GROM']
    return SimNAMES, Chargenames, PotNAMES

def crystal_matrix_to_lattice_parameters(crystal_matrix):
    """
    This function takes any strained crystal lattice matrix and return the lattice parameters

    **Required Inputs
    crystal_matrix = crystal lattice matrix ([[Vxx,Vxy,Vxz],
                                              [Vyx,Vyy,Vyz],
                                              [Vzx,Vzy,Vzz]])
    """
    # Computing lattice parameters
    a = np.linalg.norm(crystal_matrix[:, 0])
    b = np.linalg.norm(crystal_matrix[:, 1])
    c = np.linalg.norm(crystal_matrix[:, 2])

    gamma = np.arccos(np.dot(np.squeeze(np.asarray(crystal_matrix[:, 0])), np.squeeze(np.asarray(crystal_matrix[:, 1])))
                      / (a * b)) * 180. / np.pi
    alpha = np.arccos(np.dot(np.squeeze(np.asarray(crystal_matrix[:, 1])), np.squeeze(np.asarray(crystal_matrix[:, 2])))
                      / (b * c)) * 180. / np.pi
    beta = np.arccos(np.dot(np.squeeze(np.asarray(crystal_matrix[:, 2])), np.squeeze(np.asarray(crystal_matrix[:, 0])))
                     / (c * a)) * 180. / np.pi

    # Creating an array of lattice parameters
    lattice_parameters = np.array([a, b, c, alpha, beta, gamma])
    return lattice_parameters

def dGvsT(plot_out=False, Temperatures=np.array([100,200,300]), Temperatures_unsampled=[], Pressure=1, Molecules=72, molecule='benzene',
          Independent=0, potential='oplsaa', ignoreframes=200, includeframes=100000,
          simulation='gromacs', directory='', ensemble='NVT', spacing=1, hinge='DefaultHinge', phase='solid',
          Polymorphs=['p1', 'p2', 'p3'], refT=200, refdG=[0.000, 0.185, 0.306], refddG=[0.000, 0.019, 0.019],
          refdU=[0.000, 0.267, 0.240], absolutedU=[-5.624, -5.362, -5.380], output_directory='output'):
#NSA: Is this needed?
    Colors = ['b', 'g', 'r', 'm', 'c', 'y', 'k', 'g', 'g', 'g', 'g', 'g', 'g', 'g', 'g', 'g', 'g', 'g', 'g']

    if Independent == 0:
        Independent = Molecules

    # Hard set from old dictionary funciton
    refPot = 0
    ExtraPressures = []

    if np.all(Temperatures_unsampled == None):
        Temperatures_unsampled = np.array([])

    if (not np.any(Temperatures == refT)) and (not np.any(Temperatures_unsampled == refT)):
        Temperatures_unsampled = np.append(Temperatures_unsampled, refT)

    Temperatures = np.sort(np.append(Temperatures, Temperatures_unsampled))
    Pressures = np.ones(len(Temperatures), int)
    Pressures[len(Pressures) - len(ExtraPressures): len(Pressures)] = ExtraPressures
    Potentials = [potential]

    if (plot_out):
#        import matplotlib.cm as cm
#        from matplotlib.font_manager import FontProperties as FP
        font = {'family': 'normal',
                'weight': 'normal',
                'size': 14}
        matplotlib.rc('font', **font)

    # =============================================================================================
    # ENSURE THAT USER INPUTS ARE SENSIBLE
    # =============================================================================================
    # Pressure
    if Pressure < 0:
        print("Invalid Pressure: " + str(Pressure))
        sys.exit()

    # ENSEMBLE
    if ensemble != "NVE" and ensemble != "NVT" and ensemble != "NPT":
        print("Invalid Ensemble")
        print("Supported Ensembles: NVE NVT NPT")
        sys.exit()

    # =============================================================================================
    # FORMAT INPUTS
    # =============================================================================================
    # TEMPERATURE
    refk = -1
    for k, temp in enumerate(Temperatures):
        if temp == refT and refk == -1:
            refk = k + refPot * len(Temperatures)

    # =============================================================================================
    # READ IN RAW DATA
    # =============================================================================================
    # Constants.
    kB = 1.3806488e-23 * 6.0221413e23 / (1000.0 * 4.184)  # Boltzmann constant in kcal/mol/K

    # Parameters
    # How many states?
    K = len(Potentials) * len(Temperatures)

    #  maximum number of snapshots/simulation (could make this automated) - doesn't matter, as long as it's long enough.
    N_max = 5000

    # beta factor for the different temperatures
    beta_k = 1.0 / (kB * Temperatures)
    beta_k = np.tile(beta_k, (1, len(Potentials)))[0]

    # Conversion from kJ to kcal
    kJ_to_kcal = 0.2390057

    # This is the sampling efficiency for each potential in each combination of potentials
    Efficiency = np.zeros(K, float)

    # Allocate storage for simulation data
    # u_pklnT is the complete matrix of all energy data. 'p' is the polymorph, 'k' is the sampled state, 'l' is the
    #    evaluated state, 'n' is the sample number,  and T is the energy term
    #u_pklnT = np.zeros([len(Polymorphs), K, K, N_max, 20])

    # N_k[k] is the total number of snapshots from alchemical state k
    N_k = np.zeros(K, np.int32)

    # Terms_l is the list of all energy terms for the 'l' state
    Terms_l = []

    # dA[p,i,k] is the free energy between potential 0 and state k for spacing i in polymorph p
    dA = np.zeros([len(Polymorphs), spacing + 1, K], float)

    # ddA[p,i,k] is the uncertainty in the free energy between potential 0 and state k for spacing i in polymorph p
    ddA = np.zeros([len(Polymorphs), spacing + 1, K], float)

    # dG[p,i,t] is the free energy between polymorph 1 and polymorph p for spacing i and temperature t
    dG = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])

    # ddG[p,i,t] is the uncertanity in the free energy between polymorph 1 and polymorph p for spacing i and temperature t
    ddG = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])

    # dS[p,i,t] is the relative entropy between polymorph 1 and polymorph p for spacing i and temperature t
    dS = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])

    # ddS[p,i,t] is the uncertanity in the relative entropy between polymorph 1 and polymorph p for spacing i and temperature t
    ddS = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])

    dS_mbar = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])
    ddS_mbar = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])
    dH_mbar = np.zeros([len(Polymorphs), spacing + 1, len(Temperatures)])

    # O_pij[p,i,j] is the overlap within polymorph p between temperature state i and temperature state j
    O_pij = np.zeros([len(Polymorphs), len(Temperatures), len(Temperatures)])
    dU = np.zeros([len(Polymorphs), len(Temperatures)])
    ddU = np.zeros([len(Polymorphs), len(Temperatures)])

    # u_kln[k,l,n] is the reduced potential energy of configuration n from potential k in potential l
    u_kln = np.zeros([K, K, N_max], np.float64)

    # V_pkn is the volume of configuration n of polymorph p at temperature k
    V_pkn = np.zeros([len(Polymorphs), len(Temperatures), N_max], float)

    # V_avg is the average volume of polymorph p at temperature k
    V_avg = np.zeros([len(Polymorphs), len(Temperatures)], float)

    # ddV_avg is the standard deviation of the volume of polymorph p at temperature k
    ddV_avg = np.zeros([len(Polymorphs), len(Temperatures)], float)

    # C_pkn is the lattice tensor of the polymorph p at temperature k
    box_place = np.matrix([[0, 0], [1, 1], [2, 2], [0, 1], [0, 2], [1, 2]])
    C_pkn = np.zeros([len(Polymorphs), len(Temperatures), N_max, 3, 3], float)

    # h_avg is the average lattice parameters of polymorph p at temperature k
    h_avg = np.zeros([len(Polymorphs), len(Temperatures), 6], float)

    # dh is the standard deviation of the lattice parameters of polymorph p at temperature k
    dh = np.zeros([len(Polymorphs), len(Temperatures), 6], float)

    # Cycle through all polymorphs
    for p, polymorph in enumerate(Polymorphs):
        # Cycle through all sampled potentials
        for i, potential_k in enumerate(Potentials):
            count = 0
            for t in range(len(Temperatures)):
                k = len(Temperatures) * i + t
                # Cycle through all evaluated potentials
                for j, potential_l in enumerate(Potentials):
                    l = len(Temperatures) * j

                    dirpath = polymorph + '/temperature/' + str(count) + '/'
                    if os.path.isfile(dirpath + 'PROD.edr') and (Temperatures[t] not in Temperatures_unsampled):
                        count += 1
                        print("loading " + dirpath + 'PROD.edr')
                        all_energy = panedr.edr_to_df(dirpath + 'PROD.edr')
                        if len(all_energy['Potential'].values) > N_max:
                            [start_production, _, _] = timeseries.detectEquilibration(all_energy['Potential'].values[::10])
                            start_production *= 10
                        else:
                            [start_production, _, _] = timeseries.detectEquilibration(all_energy['Potential'].values)

                        # Now read in the lattice tensor and average them
                        if 'Box-XX' in list(all_energy):
                            box_letters = ['XX', 'YY', 'ZZ', 'YX', 'ZX', 'ZY']
                        else:
                            box_letters = ['X', 'Y', 'Z']

                        for b in range(len(box_letters)):
                            if len(all_energy['Potential'].values) > N_max:
                                [hold,_,_] = timeseries.detectEquilibration(all_energy['Box-' + box_letters[b]].values[::10])
                                hold *= 10
                            else:
                                [hold,_,_] = timeseries.detectEquilibration(all_energy['Box-' + box_letters[b]].values)

                            if hold > start_production:
                                start_production = hold

                        if len(all_energy['Total Energy'].values[start_production:]) > N_max:
                            start_production = len(all_energy['Total Energy'].values) - N_max

                        # Setting the end point of the simulation
                        N = len(all_energy['Total Energy'].values[start_production:])
                        N_k[k] = N

                        u_kln[k, l, :N] = all_energy['Potential'].values[start_production:]

                        # Now set these energies over all temperatures
                        u_kln[k, l:(l + len(Temperatures)), :N] = u_kln[k, l, :N]

                        # Now read in the volumes and average them
                        V_pkn[p, t, :N] = all_energy['Volume'].values[start_production:]
                        V_avg[p, t] = np.average(V_pkn[p, t, :N]) / float(Independent)
                        ddV_avg[p, t] = np.std(V_pkn[p, t, :N]) / float(Independent)

                        # Making the lattice tensor all the correct sign with time
                        if count == 1:
                            sign = np.sign(md.load(dirpath + 'pre_EQ.gro').unitcell_vectors[0].T)
                            #sign = np.sign(md.load(polymorph + '/temperature/0/ANNEAL.gro').unitcell_vectors[0].T)
                            for s in range(3):
                                for j in range(3):
                                    if sign[s, j] == 0.:
                                        # Correcting for the sign of the lattice parameters
                                        sign[s, j] = 1.

                        for b in range(len(box_letters)):
                            C_pkn[p, t, :N, box_place[b, 0], box_place[b, 1]] = np.absolute(all_energy['Box-' + box_letters[b]].values[start_production:]) * \
                                    sign[box_place[b, 0], box_place[b, 1]] * 10
                        C_avg = np.average(C_pkn[p, t, :N], axis=0)
                        dC = np.std(C_pkn[p, t, :N], axis=0)
                        h_avg[p, t] = crystal_matrix_to_lattice_parameters(C_avg)
                        dh[p, t] = np.absolute(crystal_matrix_to_lattice_parameters(C_avg + dC) - h_avg[p, t])
                    else:
                        N_k[k] = 0
                        V_avg[p, t] = np.nan
                        ddV_avg[p, t] = np.nan
                        h_avg[p, t] = np.nan
                        dh[p, t] = np.nan

        print("Start1")
        # Convert all units to kcal
        #u_pklnT[p, :, :, :] *= kJ_to_kcal
        u_kln *= kJ_to_kcal

        print("Start2")
        # If this was already in kcal or already fully independent, revert
        for j in range(len(Potentials)):
            if Potentials[j][:6] == "amoeba":
                #u_pklnT[p, :, j * len(Temperatures):(j + 1) * len(Temperatures), :, :] /= kJ_to_kcal
                u_kln[:, j * len(Temperatures):(j + 1) * len(Temperatures), :] /= kJ_to_kcal

        print("Start3")
        # Remove dependent molecules
        for j in range(len(Potentials)):
            if Potentials[j][:6] != "amoeba":
                #u_pklnT[p, :, j * len(Temperatures):(j + 1) * len(Temperatures), :, :] *= float(Independent) / Molecules
                u_kln[:, j * len(Temperatures):(j + 1) * len(Temperatures), :] *= float(Independent) / Molecules

        print("Start4")
        # Now average together the energies and volumes at each state
        for t in range(len(Temperatures)):
            dU[p, t] = np.average(u_kln[t, t, :N_k[t]]) / float(Independent)
            ddU[p, t] = np.std(u_kln[t, t, :N_k[t]]) / N_k[t] ** 0.5 / float(Independent)

        print("Start5")
        # convert to nondimensional units from kcal/mol
        for k, beta in enumerate(beta_k):
            u_kln[:, k, :] *= beta

        u_kln_save = u_kln.copy()
        N_k_save = N_k.copy()
        print("End!")

        print("Number of retained samples")
        print(N_k)

        # Now create the full N_k matrix including the roll-backs as well as the free energy container
        # N_k_matrix[i,k] is the total number of snapshots from alchemical state k using in spacing i
        N_k_matrix = np.zeros([spacing + 1, K], np.int32)
        for i in range(spacing + 1):
            N_k_matrix[i, :] = N_k_save.copy()
            N_k_matrix[i, 0: len(Temperatures)] = N_k_matrix[i, 0:len(Temperatures)] * float(i) / float(spacing)

        # =============================================================================================
        # COMPUTE FREE ENERGY DIFFERENCE USING MBAR FOR EACH SPACING
        # =============================================================================================
        for i in range(spacing+1):
            if i == 0 and len(Potentials) == 1:
                continue
            # Initialize MBAR.
            print("Running MBAR...")

            # generate the weights of each of the umbrella set
            mbar = pymbar.MBAR(u_kln, N_k_matrix[i, :], verbose=True)
            print("MBAR Converged...")

            hold = mbar.computeEffectiveSampleNumber(verbose=True)
            print(hold)

            # extract self-consistent weights and uncertainties
            (df_i, ddf_i, theta_i) = mbar.getFreeEnergyDifferences()

            # extract entropy
            [_, _, Delta_u_ij, _, Delta_s_ij, dDelta_s_ij] = mbar.computeEntropyAndEnthalpy()
            print("Free Energies Optained...")

            # Store the dimensionless results in the dA container
            dA[p, i, :] = df_i[refk]
            dH_mbar[p, i, :] = Delta_u_ij[0]
            dS_mbar[p, i, :] = Delta_s_ij[0]
            ddS_mbar[p, i, :] = dDelta_s_ij[0]
            print(dA)

        # =============================================================================================
        # COMPUTE UNCERTAINTY USING MBAR
        # =============================================================================================
        g_k = np.zeros([K])
        for i in range(spacing + 1):
            if i == 0 and len(Potentials) == 1:
                continue

            for k in range(K):
                # subsample correlated data - for now, use energy from current state
                if N_k_matrix[i, k] > 0:
                    print(N_k_matrix[i, k])
                    g_k[k] = timeseries.statisticalInefficiency(u_kln_save[k, k, 0:100])
                    print("Correlation time for phase (%s), sampled state %d is %10.3f" % (phase, k, g_k[k]))

                    # subsample the data to get statistically uncorrelated data
                    indices = np.array(timeseries.subsampleCorrelatedData(u_kln_save[k, k, 0:N_k_matrix[i, k]],
                                                                          g=g_k[k]))
                    N_k_matrix[i, k] = len(indices)
                    u_kln[k, :, 0:N_k_matrix[i, k]] = u_kln_save[k, :, indices].transpose()  # not sure why we have to transpose

            print("Number of retained samples")
            print(N_k)

            print("Running MBAR...")

            # generate the weights of each state
            mbar = pymbar.MBAR(u_kln, N_k_matrix[i, :], verbose=True)
            print("MBAR Converged...")

            # extract self-consistent weights and uncertainties
            (df_u, ddf_u, theta_u) = mbar.getFreeEnergyDifferences()

            # calculate the overlap it necessary
            if len(Temperatures) == 2:
                O_pij[p, :, :] = mbar.computeOverlap()[2]

            # testing
            weights_in_gromos = np.zeros(K, float)
            for k in range(K):
                w = np.exp(mbar.Log_W_nk[:, k])
                print("max weight in state %d is %12.7f" % (k, np.max(w)))
                neff = 1 / np.sum(w ** 2)

                print("Effective number of sample in state %d is %10.3f" % (k, neff))
                print("Efficiency for state %d is %d/%d = %10.4f" % (k, neff, len(w), neff / len(w)))
                Efficiency[k] = neff / len(w)  # Store the efficiency
                w_0 = np.exp(mbar.Log_W_nk[:, 0])  # Weights in gromos
                initial_configs = np.sum(N_k[0:k])
                final_configs = np.sum(N_k[0:k + 1])

                print("Total weight in gromos " + str(np.sum(w_0[initial_configs:final_configs])))
                weights_in_gromos[k] = np.sum(w_0[initial_configs:final_configs])

            # Write out free energy differences
            print("Free Energy Difference (in units of kcal/mol)")
            for k in range(K):
                print("%8.3f %8.3f" % (-df_i[k, 0], ddf_u[k, 0]))

            # Store the dimensionless results in the ddA container
            ddA[p, i, :] = ddf_u[refk]

    # Check the overlap it necessary
    if len(Temperatures) == 2:
        print("Overlap:")
        print(O_pij)
        pdb.set_trace()

    # =============================================================================================
    # FINALIZE THE RELATIVE FREE ENERGY AND ENTROPY
    # =============================================================================================
    for i in range(spacing + 1):
        for t, T in enumerate(Temperatures):
            for p in range(len(Polymorphs)):
                #print('HERE!!!!!', dA[p, i, t], dA[0, i, t], (dA[p, i, t] - dA[0, i, t]), (beta_k[t] * float(Independent)), float(T),  float(refT), refdG[p])
                dG[p, i, t] = (dA[p, i, t] - dA[0, i, t]) / (beta_k[t] * float(Independent)) + float(T) / float(refT) * \
                                                                                refdG[p]
                ddG[p, i, t] = ((ddA[p, i, t] ** 2 + ddA[0, i, t] ** 2) / (beta_k[t] * float(Independent)) ** 2 +
                                float(T) / float(refT) * float(refddG[p]) ** 2) ** 0.5
                if p == 0:
                    continue
                dS[p, i, t] = (dU[p, t] - dU[0, t] - dG[p, i, t]) / float(T)
                ddS[p, i, t] = (ddU[p, t] ** 2 + ddU[p, t] ** 2 + ddG[p, i, t] ** 2) ** 0.5 / float(T)

    print("Polymorph Free Energy:")
    for p in range(len(Polymorphs)):
        print("%8.3f %8.3f" % (dG[p, spacing, len(Temperatures) - 1], ddG[p, spacing, len(Temperatures) - 1]))

    # =============================================================================================
    # PLOT THE RELATIVE FREE ENERGY VS TEMPERATURE
    # =============================================================================================

#    Temperatures2 = np.array([0] + [j for j in Temperatures])
#    Pressures2 = np.array([1] + [j for j in Pressures])
#    dG2 = np.insert(dG, 0, np.transpose(np.tile(refdU, (1, 1))), axis=2)
#    ddG2 = np.insert(ddG, 0, np.zeros([len(Polymorphs), spacing + 1]), axis=2)
#    dS2 = np.insert(dS, 0, np.zeros([len(Polymorphs), spacing + 1]), axis=2)
#    ddS2 = np.insert(ddS, 0, np.zeros([len(Polymorphs), spacing + 1]), axis=2)
    PlotPress = 1  # Pressure to plot the dGvT curve at
    Temperatures_P = Temperatures[Pressures == PlotPress]

    if plot_out == True:
        f, a = plt.subplots(1, 1)
        xlabel = 'Temperature (K)'
        ylabel = 'Relative Free Energy (kcal/mol)'

        a.errorbar(Temperatures, dG[0, 0, :], yerr=np.zeros(len(dG[0, 0, :]), float), linestyle='--', marker='.',
                   linewidth=2, alpha=0.6, color='b', label=Polymorphs[0])
        for p in range(len(Polymorphs)):
            if p == 0:
                continue
            if len(Potentials) > 1:
                a.errorbar(Temperatures, dG[p, 0, :], yerr=ddG[p, 0, :], linestyle='--', marker='.', linewidth=2,
                           alpha=0.6, color=Colors[p], label=Polymorphs[p])
            a.errorbar(Temperatures_P, dG[p, 1, Pressures == PlotPress], yerr=ddG[p, 1, Pressures == PlotPress],
                       linestyle='-', marker='.', linewidth=2, alpha=0.6, color=Colors[p], label=Polymorphs[p])

        a.set_xlabel(xlabel)
        a.set_ylabel(ylabel)

    if not os.path.isdir(output_directory):
        subprocess.call(['mkdir', output_directory])

    np.save(output_directory + '/T_' + molecule + '_' + potential, Temperatures_P)
    for p, Poly in enumerate(Polymorphs):
        np.save(output_directory + '/dGvT_' + molecule + '_' + Poly + '_' + potential, dG[p, spacing, Pressures == PlotPress])
        np.save(output_directory + '/ddGvT_' + molecule + '_' + Poly + '_' + potential, ddG[p, spacing, Pressures == PlotPress])
        if len(Potentials) > 1:
            np.save(output_directory + '/dGvT_' + molecule + '_' + Poly + '_' + potential + '_indirect', dG[p, 0, :])
            np.save(output_directory + '/ddGvT_' + molecule + '_' + Poly + '_' + potential + '_indirect', ddG[p, 0, :])
            if spacing > 1:
                np.save(output_directory + '/dGvT_' + molecule + '_' + Poly + '_' + potential + '_convergence', dG[p, :, :])
                np.save(output_directory + '/ddGvT_' + molecule + '_' + Poly + '_' + potential + '_convergence', ddG[p, :, :])
        np.save(output_directory + '/dS_' + molecule + '_' + Poly + '_' + potential, dS[p, spacing, :])
        np.save(output_directory + '/ddS_' + molecule + '_' + Poly + '_' + potential, ddS[p, spacing, :])

    # =============================================================================================
    # PLOT THE RELATIVE ENTROPY VS TEMPERATURE
    # =============================================================================================

    if plot_out == True:
        f, a = plt.subplots(1, 1)
        plt.xlabel('Temperature (K)', fontsize=20)
        plt.ylabel('Relative Entropy (kcal/molK)', fontsize=20)
        for p in range(len(Polymorphs)):
            a.errorbar(Temperatures2, dS[p, spacing, :], yerr=ddS[p, spacing, :], linestyle='--', marker='.', linewidth=2,
                       alpha=0.6, color=Colors[p], label=Polymorphs[p])

    # =============================================================================================
    # PLOT THE AVERAGE ENERGY VS TEMPERATURE
    # =============================================================================================

#    dU2 = np.insert(dU, 0, np.transpose(absolutedU), axis=1)
#    ddU2 = np.insert(ddU, 0, 0, axis=1)

    if plot_out == True:
        f, a = plt.subplots(1, 1)
        plt.xlabel('Temperature (K)', fontsize=20)
        plt.ylabel('Average Energy (kcal/mol)', fontsize=20)
        for p in range(len(Polymorphs)):
            a.errorbar(Temperatures2, dU[p, :], yerr=ddU[p, :], linestyle='--', marker='.', linewidth=2, alpha=0.6,
                       color=Colors[p], label=Polymorphs[p])
        plt.legend(loc='upper left')

    for p, Poly in enumerate(Polymorphs):
        np.save(output_directory + '/UvT_' + molecule + '_' + Poly + '_' + potential, dU[p, :])

    # =============================================================================================
    # PLOT THE AVERAGE BOX VOLUME VS TEMPERATURE
    # =============================================================================================

    if plot_out == True:
        f, a = plt.subplots(1, 1)
        plt.xlabel('Temperature (K)', fontsize=20)
        for p in range(len(Polymorphs)):
            a.errorbar(Temperatures, V_avg[p, :], yerr=ddV_avg[p, :], linestyle='--', marker='.', linewidth=2,
                       alpha=0.6, color=Colors[p], label=Polymorphs[p])

    for p, Poly in enumerate(Polymorphs):
        np.save(output_directory + '/VvT_' + molecule + '_' + Poly + '_' + potential, V_avg[p, :])
        np.save(output_directory + '/dVvT_' + molecule + '_' + Poly + '_' + potential, ddV_avg[p, :])

    # =============================================================================================
    # SAVE THE AVERAGE BOX VECTORS AND ANGLES VS TEMPERATURE
    # =============================================================================================

    for p, Poly in enumerate(Polymorphs):
        np.save(output_directory + '/hvT_' + molecule + '_' + Poly + '_' + potential, h_avg[p, :])
        np.save(output_directory + '/dhvT_' + molecule + '_' + Poly + '_' + potential, dh[p, :])

    # =============================================================================================
    # PLOT THE DIFFERENCE IN AVERAGE ENERGY VS TEMPERATURE
    # =============================================================================================

    if plot_out == True:
        f, a = plt.subplots(1, 1)
        plt.xlabel('Temperature (K)', fontsize=20)
        plt.ylabel('Average Energy Difference (kcal/mol)', fontsize=20)
        for p in range(len(Polymorphs)):
            if p == 0:
                yerror = np.zeros(len(ddU[0, :]))
            else:
                yerror = []
                for t in range(len(ddU[0, :])):
                    yerror.append(float((ddU[0, t] ** 2 + ddU[p, t] ** 2) ** 0.5))
            a.errorbar(Temperatures2, dU[p, :] - dU[0, :], yerr=yerror, linestyle='--', marker='.', linewidth=2,
                       alpha=0.6, color=Colors[p], label=Polymorphs[p])

    # Save the data for future use.
    for p, Poly in enumerate(Polymorphs):
        np.save(output_directory + '/dUvT_' + molecule + '_' + Poly + '_' + potential, dU[p, :] - dU[0, :])
        np.save(output_directory + '/ddUvT_' + molecule + '_' + Poly + '_' + potential, (ddU[p, :] ** 2 + ddU[0, :] ** 2) ** 0.5)

    if plot_out == True:
        plt.tight_layout()
        plt.show()

if __name__ == '__main__':
    # =============================================================================================
    # READ IN USER INPUTS
    # =============================================================================================
    parser = OptionParser()
    parser.add_option('-p', '--plot', dest='plot', help='Plot output (default false)', default=True, action='store_true')
#    parser.add_option('-T', dest='temp', help='Temperature', default=200)
    parser.add_option('-P', dest='Pressure', help='Pressure', default=1)
#    parser.add_option('-n', dest='polymorphs', help='Polymorphs to analyze', default='p1 p2 p3')
    parser.add_option('-N', dest='molecules', help='number of supercell molecules', default=72)
    parser.add_option('-M', dest='molecule', help='name of the molecule', default='benzene')
    parser.add_option('-I', dest='independent', help='number of independent molecules', default='Same')
    parser.add_option('-u', dest='potential', help='potential to create the phase diagram in', default='oplsaa')
    parser.add_option('-i', dest='ignoreframes', help='Initial frames to ignore', default=200)
    parser.add_option('-j', dest='includeframes', help='Number of frames to include', default=100000)
    parser.add_option('-z', dest='simulation', help='The simulation package that was used', default='gromacs')
    parser.add_option('-d', dest='directory', help='Parent directory of the reweight directories', default='None')
    parser.add_option('-E', dest='ensemble', help='Simulation Ensemble', default='NVT')
    parser.add_option('-s', dest='spacing', help='spacing for the plot rolling back the state 0 sampling', default=1)
    parser.add_option('-H', '--hinge', dest='hinge', help='Optional string at end of jobs', default='DefaultHinge')

    (options, args) = parser.parse_args()
    plot_out = options.plot
    Temp = options.temp
    Pressure = int(options.Pressure)
    Molecules = int(options.molecules)
    molecule = options.molecule
    if options.independent == 'Same':
        Independent = Molecules
    else:
        Independent = int(options.independent)
    potential = options.potential
    ignoreframes = int(options.ignoreframes)
    includeframes = int(options.includeframes)
    simulation = options.simulation
    directory = options.directory
    if directory == 'None':
        directory = ''
    ensemble = options.ensemble
    spacing = int(options.spacing)
    hinge = options.hinge
    phase = "solid"

    Polymorphs, refT, refdG, refddG, refdU, absolutedU = old_systems_dictionary(potential, molecule)
    dGvsT(plot_out=plot_out, Temperatures=Temp, Temperatures_unsampled=[], Pressure=Pressure, Molecules=Molecules, molecule=molecule,
          Independent=Independent, potential=potential, ignoreframes=ignoreframes, includeframes=includeframes,
          simulation=simulation, directory=directory, ensemble=ensemble, spacing=spacing, hinge=hinge, phase=phase,
          Polymorphs=Polymorphs, refT=refT, refdG=refdG, refddG=refddG, refdU=refdU, absolutedU=absolutedU)


""" Parameters module. """

#  ISC License
#
#  Copyright (c) 2020–2022, Paul Wilhelm <anfrage@paulwilhelm.de>
#
#  Permission to use, copy, modify, and/or distribute this software for any
#  purpose with or without fee is hereby granted, provided that the above
#  copyright notice and this permission notice appear in all copies.
#
#  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
#  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
#  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
#  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
#  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
#  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

from __future__ import annotations
from typing import Callable
import numpy as np
from numba import jit, prange
from magneticalc.Backend_Types import get_jit_enabled
from magneticalc.ConditionalDecorator import ConditionalDecorator
from magneticalc.Constants import Constants
from magneticalc.Debug import Debug
from magneticalc.Field_Types import FIELD_TYPE_A, FIELD_TYPE_B
from magneticalc.Metric import Metric
from magneticalc.Validatable import Validatable, require_valid, validator

class Parameters(Validatable):
    """ Parameters class. """

    def __init__(self) -> None:
        """
        Initializes parameters class.
        """
        Validatable.__init__(self)
        Debug(self, ": Init", init=True)

        self._energy: float = 0.0
        self._self_inductance: float = 0.0
        self._magnetic_dipole_moment: float = 0.0

    def set(self, *args, **kwargs):
        """
        Sets the parameters
        """

    # ------------------------------------------------------------------------------------------------------------------

    def _get_squared_field(
            self,
            sampling_volume: SamplingVolume,  # type: ignore
            field: Field  # type: ignore
    ) -> float:
        """
        Returns the "squared" field scalar.

        @param sampling_volume: SamplingVolume
        @param field: B-field
        @return: Float
        """
        return self._get_squared_field_worker(sampling_volume.permeabilities, field.vectors)

    @staticmethod
    @ConditionalDecorator(get_jit_enabled(), jit, nopython=True, parallel=True)
    def _get_squared_field_worker(sampling_volume_permeabilities: np.ndarray, field_vectors: np.ndarray) -> float:
        """
        Returns the "squared" field scalar.

        @param sampling_volume_permeabilities: Ordered list of sampling volume's relative permeabilities µ_r
        @param field_vectors: Ordered list of 3D vectors (B-field)
        @return: Float
        """
        squared = 0
        for i in prange(len(field_vectors)):
            squared += np.dot(field_vectors[i], field_vectors[i] / sampling_volume_permeabilities[i])
        return squared

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def _get_magnetic_dipole_moment(
            self, wire: Wire,  # type: ignore
            length_scale: float
    ) -> float:
        """
        Returns the magnetic dipole moment scalar.

        @param wire: Wire
        @param length_scale: Length scale (m)
        @return: Float
        """
        elements_center = np.array([element[0] for element in wire.elements])
        elements_direction = np.array([element[1] for element in wire.elements])
        vector = self._get_magnetic_dipole_moment_worker(elements_center, elements_direction, length_scale)
        return np.abs(wire.dc * np.linalg.norm(vector) / 2)

    @staticmethod
    @ConditionalDecorator(get_jit_enabled(), jit, nopython=True, parallel=True)
    def _get_magnetic_dipole_moment_worker(
            elements_center: np.ndarray,
            elements_direction: np.ndarray,
            length_scale: float
    ):
        """
        Returns the (unscaled) magnetic dipole moment vector.

        @param elements_center: Current element centers
        @param elements_direction: Current element directions
        @param length_scale: Length scale (m)
        @return: Magnetic dipole moment vector
        """
        squared = np.zeros(3)
        for i in prange(len(elements_center)):
            squared += np.cross(elements_center[i] * length_scale, elements_direction[i] * length_scale)
        return squared

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    @validator
    def recalculate(
            self,
            wire: Wire,  # type: ignore
            sampling_volume: SamplingVolume,  # type: ignore
            field: Field,  # type: ignore
            progress_callback: Callable
    ) -> bool:
        """
        Recalculates parameters.

        @param wire: Wire
        @param sampling_volume: SamplingVolume
        @param field: Field
        @param progress_callback: Progress callback
        @return: True (currently non-interruptable)
        """
        Debug(self, ".recalculate()")

        progress_callback(0)

        self._magnetic_dipole_moment = self._get_magnetic_dipole_moment(wire, Metric.LengthScale)

        progress_callback(33)

        if field.type == FIELD_TYPE_A:

            pass

        elif field.type == FIELD_TYPE_B:

            dV = (Metric.LengthScale / sampling_volume.resolution) ** 3  # Sampling volume element
            self._energy = self._get_squared_field(sampling_volume, field) * dV / Constants.mu_0

            progress_callback(66)

            self._self_inductance = self._energy / np.square(wire.dc)

        progress_callback(100)

        return True

    @property
    @require_valid
    def energy(self) -> float:
        """
        Returns calculated energy.

        @return: Float
        """
        return self._energy

    @property
    @require_valid
    def self_inductance(self) -> float:
        """
        Returns calculated self-inductance.

        @return: Float
        """
        return self._self_inductance

    @property
    @require_valid
    def magnetic_dipole_moment(self) -> float:
        """
        Returns calculated magnetic dipole moment.

        @return: Float
        """
        return self._magnetic_dipole_moment


# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Some code is from https://github.com/princeton-vl/pose-ae-train/blob/454d4ba113bbb9775d4dc259ef5e6c07c2ceed54/utils/group.py
# Written by Bin Xiao (leoxiaobin@gmail.com)
# Modified by Bowen Cheng (bcheng9@illinois.edu)
# ------------------------------------------------------------------------------
"""
############################################################################
将检测分组为个体实例姿态。为此, 网络将在每个关节的像素位置生成一个标记。
换句话说, 每个关节热图都有相应的 "tag标签" 热图。
因此, 如果有m个人体关节需要预测, 那么网络将输出总共2m通道, m用于检测和m用于分组。
为了将检测解析成个体实例,
我们使用非最大抑制来获取每个关节的峰值检测, 并在相同的像素位置检索其相应的tag标记。
然后, 我们通过比较检测的标记值和匹配足够接近的标记值, 对整个人体关节点检测进行分组。
而后一组检测形成了一个人的姿势估计。
############################################################################
为了产生最后一个最终预测集合，我们对每个关节点进行迭代。
迭代的顺序是先考虑头和躯干而后逐渐像肢体移动。
我们从第一个关节点开始，NMS后得到每个超过阈值的关节点。
这些关键点组组成了我们待检测人的最初候选池。
我们之后考虑检测后续的关节点。我们将后续关节点与我们现有关节点候选池相比，得到其最优匹配。
只有当两个标记相差小于一个特别的阈值时才会被匹配。
除此之外，我们因此基于标记距离和检测分数的最大匹配。
如果没有任何检测额被匹配，则被视为一个新的人体实例。这种情况可能是因为某个特定人体只有手或脚可见。
我们在每类关节点上循环匹配，直到每个检测都被匹配到一个人。
##########################################################################################
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from munkres import Munkres
import numpy as np
import torch

def py_max_match(scores):
    m = Munkres()   # 匈牙利算法
    """
    ##################################################################
    Compute the indexes for the lowest-cost pairings between rows and
    columns in the database. Returns a list of `(row, column)` tuples
    that can be used to traverse the matrix.
    ##################################################################
    """
    tmp = m.compute(scores)
    tmp = np.array(tmp).astype(np.int32)
    return tmp

"""
##########################
# 通过tag值，对节点进行匹配
##########################
# 一张图片检测出了几个人，最终有几个人，就有几个tag键，最后在
            # grouped_keys中就会有几个不同的值，每个人对应一个值，
            # grouped_keys: [ tag1,tag2,...,tag_?? ]
            # joints_dict: { tag1: array(17*5) , tag2: array(17*5), .... ,tag_??: array(17*5) }    ，
            # 最终有几个人，就有几个tag键
            # grouped_key中的值和joints_dict中的键一致。
            # grouped_tags形如:[array([ 3.0320523, -2.9441261] ), array([-3.2409573,  3.267928 ] ), array([-0.46144685,  0.43847117] ), array([-6.7198505,  6.6868863] )]
            # tag_dict形如:
            #            {3.0482664: [array([ 3.0482664, -2.942074 ] ), array([ 3.0499892, -2.9416664] ), array([ 3.04908 , -2.946251] ), array([ 3.0546198, -2.9452639] ), array([ 3.0284576, -2.9707832] ), array([ 3.011262 , -2.9344423] ), array([ 3.0276155, -2.9461455] ), array([ 3.022631, -2.930122] ), array([ 3.0320063, -2.947176 ] ), array([ 3.0340366, -2.9682217] ), array([ 3.0116858, -2.9123354] ), array([ 3.0149767, -2.9450297] ), array([ 3.0390344, -2.9229352] )],
            #            -3.204033: [array([-3.204033 ,  3.2672586] ), array([-3.224389 ,  3.2966106] ), array([-3.2322555,  3.254828 ] ), array([-3.2290933,  3.2748113] ), array([-3.2254848,  3.290883 ] ), array([-3.2734041,  3.2389045] ), array([-3.262566 ,  3.2427683] ), array([-3.270686,  3.308473] ), array([-3.240734 ,  3.2396798] ), array([-3.22856  ,  3.2748213] ), array([-3.2633042,  3.247148 ] ), array([-3.2369757,  3.2789457] ), array([-3.2354915,  3.1973739] )],
            #            -0.48780698: [array([-0.48780698,  0.4803399 ] ), array([-0.4915051 ,  0.48355848] ), array([-0.47469807,  0.46150476] ), array([-0.43962345,  0.45701584] ), array([-0.47148287,  0.45973817] ), array([-0.38883165,  0.44795612] ), array([-0.39312   ,  0.29450473] ), array([-0.49791056,  0.4785707 ] ), array([-0.4475416 ,  0.36834407] ), array([-0.47388119,  0.49372977] ), array([-0.47124237,  0.3472029 ] ), array([-0.49971843,  0.48918843] ), array([-0.49369287,  0.46173176] )],
            #            -6.7378864: [array([-6.7378864,  6.6715736] ), array([-6.738204,  6.720743] ), array([-6.7024045,  6.702105 ] ), array([-6.674477 ,  6.6906085] ), array([-6.7082214,  6.6431236] ), array([-6.8019643,  6.791915 ] ), array([-6.6872525,  6.6337585] ), array([-6.708395 ,  6.6412625] ), array([-6.686763,  6.692321] )]
            #            }
##################################################
"""

def match_by_tag(inp, params):
    assert isinstance(params, Params), 'params should be class Params()'
    """
    # zip后的inp，使得：
    # tag_k.shape=(17,30,2)
    # loc_k.shape=(17,30,2)
    # val_k.shape=(17,30)
    """
    tag_k, loc_k, val_k = inp

    default_ = np.zeros((params.num_joints, 3 + tag_k.shape[2]))   # shape=(17,5)

    joint_dict = {}
    tag_dict = {}
    for i in range(params.num_joints): # 17
        idx = params.joint_order[i]
        """
        # 17类中的第idx类节点
        """
        # tags.shape=(30,2)
        tags = tag_k[idx]
        # joints.shape=(30,2+1+2)  , 二维坐标 + value值 + tag值
        joints = np.concatenate(
            (loc_k[idx], val_k[idx, :, None], tags), 1
        )
        # 从30个节点中选择val值大于阈值的节点,
        # 这些关键点组组成了我们待检测人的最初候选池.
        # mask.shape=(30)
        mask = joints[:, 2] > params.detection_threshold

        """
        ##############################################
        # 从该类节点的30个中，
        # 挑选出val大于阈值的节点对应的tags和joints
        # 假设有m个节点满足阈值条件，则
        # tags.shape=(k,2)
        # joints.shape=(k,5)
        ##############################################
        """
        # tags.shape=(k,2)
        tags = tags[mask]
        # joints.shape=(k,5)  , 二维坐标 + value值 + tag值
        print('joints.shape:{}'.format(joints.shape))

        if joints.shape[0] == 0:
            continue

        """
        注意：注意：注意：注意：
        这些键tag都是第一次出现的值（所以几乎都是第一类节点中的tag值，不过如果某些人的第一类节点可能被遮挡，这些人的tag键开始时就没有出现）
        """
        # 当此时访问的是第一类节点
        if i == 0 or len(joint_dict) == 0:
            for tag, joint in zip(tags, joints):   # 逐一访问第一类节点的30个中的k=m个
                # tag.shape=2 , joint.shape=5
                # 以该点的tag值作为字典的键
                # 共生成m个键值对
                key = tag[0]
                # Python字典setdefault()方法和get()方法类似,如果键不已经存在于字典中,将会添加键并将值设为默认值。

                joint_dict.setdefault(key, np.copy(default_))[idx] = joint
                tag_dict[key] = [tag]

        # 当该类节点不是第一类节点时
        # 假设此时访问第二类节点
        else:
            # grouped_keys和grouped_tags为之前的值
            # grouped_keys = [ tag_1 , tag_2 , ... , tag_?? ]
            grouped_keys = list(joint_dict.keys())[:params.max_num_people]

            # grouped_tags = [ array([a_1,b_1]) , array([a_2,b_2]) , ... , array([a_????,b_????]) ]
            """
            tag_dict: 不同的键tag代表不同一个人，而在一个键tag中，保存着该人存在的节点的真实tag值，
            因为一个人的节点最多有17个，所以一个键tag最多会有17个节点的tag值,最少一个tag值。
            """
            grouped_tags = [np.mean(tag_dict[i], axis=0) for i in grouped_keys]

            # False
            if params.ignore_too_much \
               and len(grouped_keys) == params.max_num_people:
                continue

            """
            ##########################################################
            # joints为当前节点，假设有k=n个候选点，则joints.shape=(n,5)，
            # grouped_tags为之前的tags值。np.array(grouped_tags) = array([[a_1,b_1],[a_2,b_2],...,[a_m,b_m]])  , shape=(m,2)
            ##################################################################################################################
            """
            """
            $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$
            # 计算了“当前类别的n个节点”分别与“之前的m个节点”差值diff=(n,1,2)-(1,m,2),再求L2范数
            # 得到diff.shape=(n,m)
            $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$
            注意：
            注意：
            注意：
            当前类别的节点是与到目前为止已存在的人的代表节点的tag(即：已存在的字典的键)。
            注意：
            注意：
            注意：
            当前类别的节点是与到目前为止已存在的人的代表节点的tag(即：已存在的字典的键)。
            注意：
            """
            diff = joints[:, None, 3:] - np.array(grouped_tags)[None, :, :]  # diff.shape=(n,m,2)
            diff_normed = np.linalg.norm(diff, ord=2, axis=2)    # diff_normed.shape=(n,m)

            diff_saved = np.copy(diff_normed)

            if params.use_detection_val: # True
                diff_normed = np.round(diff_normed) * 100 - joints[:, 2:3]

            num_added = diff.shape[0]    # 当前类节点的 k2 个值
            num_grouped = diff.shape[1]  # 之前类节点的 k1 个值

            # 如果 “当前类节点数n” 比 “之前类别节点m” 大。
            if num_added > num_grouped:
                # np.concatenate(([n,m],[n,n-m]),axis=1) ----> diff_normed.shape=( n , n )
                diff_normed = np.concatenate(
                    (
                        diff_normed,
                        np.zeros((num_added, num_added-num_grouped))+1e10
                    ),
                    axis=1
                )

            """
            ###########################################
            # 运用匈牙利算法寻找相邻节点的最优匹配。
            # 获取相邻节点的匹配对,pair.shape=(n,2)
            ###########################################
            """
            pairs = py_max_match(diff_normed)   # 运用匈牙利算法

            for row, col in pairs:
                if (row < num_added and col < num_grouped and diff_saved[row][col] < params.tag_threshold):
                    """
                    # 注意：
                    # 注意：
                    # 注意：
                    # 注意：
                    # 当前的tag值与之前的键tag相匹配。
                    # 说明当前值为tag的节点属于已经存在的一个人。
                    # 注意：
                    # 注意：
                    """
                    key = grouped_keys[col]       # grouped_ksys[col]是之前已检测到的人的键tag   , 其与当前的tag是最佳匹配。
                    joint_dict[key][idx] = joints[row]  # joints[row]是当前的n个节点的第row个节点的joint ，将其添加到joint_dict[前tag][第idx类节点]
                    tag_dict[key].append(tags[row])     # tags[row]是当前的n个节点的第row个节点的tag ，将其添加到tag_dict[前tag]
                else:
                    """
                    # 注意：
                    # 注意：
                    # 注意：
                    # 当当前类别的某一节点的tag与已存在的人的代表tag(即：已存在的字典键tag)都不匹配时，
                    # 说明：当前出现的不能匹配tag值表明了一个新的人体出现，
                    # 所以，需要更新旧的字典，添加新的键tag，并将该节点(即该人体)的joints和tagsdian加入该键值中。
                    # 注意：
                    # 注意：
                    """
                    key = tags[row][0]
                    joint_dict.setdefault(key, np.copy(default_))[idx] = \
                        joints[row]
                    tag_dict[key] = [tags[row]]

            """
            ####################################################################################################################################
            # 最终，
            # 得到该幅图片的所有检测到的一个一个单独的人体肢体节点集合。
            # joint_dict:{tag1:array(x1,5),tag2:array(x2,5),...,tag_i:array(xi,5)}
            # 其中，有几个人体就有几个tag键，并且每个人体出现的关节点数也是不相同的，即每个"x"的数值是不定的，x在(1-->17)之间。
            # 注意：
            # 注意：
            # 注意：这里的idx不是一个连续的数值，每一类的节点都有对应的idx，并且初始化的数组default_=np.zeros(17,5)！！！！！！！！！！！！！！！
            # 注意：因此，最后的joint_dict中的数组array是有全零行的；最终的每个人体的数组shape=(17,5)！！！！！！！！！！！！！！！！！！！！！！
            # 注意：
            # 注意：
            ------------------------------------------------------------------------------------------------------------------------------------
            # 其中，tag_dict和joint_dict是相对应的，joint_dict包含tag_dict ,
            # tag_dict:{tag1:[array(tag11),..,array(tag_x1)],tag2:[array(tag22),..,array(tag_x2)],....,tag_i:[array(tagii),..,array(tag_xi)]}
            #####################################################################################################################################
            """

    ans = np.array([joint_dict[i] for i in joint_dict]).astype(np.float32)
    # ans.shape=(人体数目，17，5)
    print('ans.shape={}'.format(ans.shape))
    return ans

class Params(object):
    def __init__(self, cfg):
        self.num_joints = cfg.DATASET.NUM_JOINTS
        self.max_num_people = cfg.DATASET.MAX_NUM_PEOPLE   # 30

        self.detection_threshold = cfg.TEST.DETECTION_THRESHOLD
        self.tag_threshold = cfg.TEST.TAG_THRESHOLD
        self.use_detection_val = cfg.TEST.USE_DETECTION_VAL
        self.ignore_too_much = cfg.TEST.IGNORE_TOO_MUCH

        if cfg.DATASET.WITH_CENTER and cfg.TEST.IGNORE_CENTER:
            self.num_joints -= 1

        if cfg.DATASET.WITH_CENTER and not cfg.TEST.IGNORE_CENTER:
            self.joint_order = [
                i-1 for i in [18, 1, 2, 3, 4, 5, 6, 7, 12, 13, 8, 9, 10, 11, 14, 15, 16, 17]
            ]
        else:
            self.joint_order = [
                i-1 for i in [1, 2, 3, 4, 5, 6, 7, 12, 13, 8, 9, 10, 11, 14, 15, 16, 17]
            ]

class HeatmapParser(object):
    def __init__(self, cfg):
        self.params = Params(cfg)
        self.tag_per_joint = cfg.MODEL.TAG_PER_JOINT
        self.pool = torch.nn.MaxPool2d(
            cfg.TEST.NMS_KERNEL, 1, cfg.TEST.NMS_PADDING
        )

    def nms(self, det):
        """
        # # 对热力图进行最大池化torch.nn.MaxPool2d(5,1,2)
        """
        maxm = self.pool(det)
        """
        # 挑选出最大值位置，maxm在该处值为1
        """
        maxm = torch.eq(maxm, det).float()
        """
        # 将热力图det中的极大值保留，去除非极大值
        """
        det = det * maxm
        return det

    def match(self, tag_k, loc_k, val_k):
        # tag_k.shape=(1,17,30,2)
        # loc_k.shape=(1,17,30,2)二维坐标(x,y)
        # val_k.shape=(1,17,30)
        match = lambda x: match_by_tag(x, self.params)
        """
        # zip() 函数用于将可迭代的对象作为参数，将对象中对应的元素打包成一个个元组，然后返回由这些元组组成的对象
        # 如果各个迭代器的元素个数不一致，则返回列表长度与最短的对象相同
        """
        """
        # match函数返回该幅图片的所有检测到的一个一个单独的人体肢体节点集合。
        # joint_dict:{tag1:array(x1,5),tag2:array(x2,5),...,tag_i:array(xi,5)}
        # 其中，有几个人体就有几个tag键，并且每个人体出现的关节点数也是不相同的，所以每个人体对应的array(17,5)中的“一些行是全零行”。
        # 接着再将其转换为numpy数组ans,ans.shape=(人体数，17，5),其中ans.shape[0]等于检测到的人体数。
        # 最后，再将其第0维的每个array(17，5)变为list列表元素。
        ################################################################

        """
        return list(map(match, zip(tag_k, loc_k, val_k)))

    """
    #############################################
    # tensor张量：热力图det.shape=(1,17,h,w)
    # tensor张量：tag图tag.shape=(1,17,h,w,2)
    #############################################
    """
    def top_k(self, det, tag):
        # det = torch.Tensor(det, requires_grad=False)
        # tag = torch.Tensor(tag, requires_grad=False)

        """
        ##################################
        # 对热力图进行NMS非极大值抑制，
        # 通过最大池化操作，找到极大值位置，
        # 保留热力图中的极大值，去除非极大值
        ##################################
        """
        det = self.nms(det)
        num_images = det.size(0)   # 1
        num_joints = det.size(1)   # 17
        h = det.size(2)     # 512
        w = det.size(3)     # 512

        """
        #######################################
        # 将每个节点的二维(h,w)热力图展开为一维
        # det.shape=(1,17,262144)
        #######################################
        """
        det = det.view(num_images, num_joints, -1)     # shape=(1,17,262144)

        """
        ###########################################################################
        # topk:沿给定dim=2维度返回输入张量input中 k=30 个最大值。
        # 如果不指定dim，则默认为input的最后一维,
        # 返回一个元组 (values,indices)，其中indices是原始输入张量input中测元素下标
        ##########################################################################
        # 得到每个节点的前30个最大值，及其对应的下标
        # val_k.shape=(1,17,30)
        # ind.shape=(1,17,30)
        # ind为“一维坐标值”
        ############################################
        """
        val_k, ind = det.topk(self.params.max_num_people, dim=2)      # 得到每个节点的前30个最大值，及其对应位置一维坐标

        """
        ###################################################
        # 将tag.shape=(1,17,512,512,2) ---> (1,17,262144,2)
        ###################################################
        """
        tag = tag.view(tag.size(0), tag.size(1), w*h, -1)

        if not self.tag_per_joint:  # not True
            tag = tag.expand(-1, self.params.num_joints, -1, -1)

        """
        ######################################################################
        # torch.gather(input,dim,index),从原tensor中获取指定dim和指定index的数据,
        # 在tag图中索取热力图中前30个人每个节点位置对应的tag值
        # tag_k.shape=(1,17,30,2)
        ######################################################################
        """
        tag_k = torch.stack(
            [
                torch.gather(tag[:, :, :, i], 2, ind)
                for i in range(tag.size(3))
            ],
            dim=3
        )

        """
        ################################
        # 将“一维坐标ind”转换为二维坐标(x,y)
        # ind.shape=(1,17,30)
        # --->
        # ind_k.shape=(1,17,30,2)
        ################################
        """
        x = ind % w
        y = (ind / w).long()    # self.long() is equivalent to self.to(torch.int64)
        ind_k = torch.stack((x, y), dim=3)   # ind_k.shape=(1,17,30,2)

        """
        ####################################
        # tag_k.shape=(1,17,30,2)
        # ind_k.shape=(1,17,30,2)二维坐标(x,y)
        # val_k.shape=(1,17,30)
        ####################################
        """
        ans = {
            'tag_k': tag_k.cpu().numpy(),
            'loc_k': ind_k.cpu().numpy(),
            'val_k': val_k.cpu().numpy()
        }

        return ans

    def adjust(self, ans, det):
        # ans为列表，图片batch_size数目len(ans)=1 ，ans=[array(num_people，17，5),]
        for batch_id, peoples in enumerate(ans):
            # 对一张图片中的所有人
            # peoples.shape=(num_people,17,5)
            for people_id, people in enumerate(peoples):
                # people.shape=(17,5)
                for joint_id, joint in enumerate(people):
                    # joint.shape=(5,)  , 二维节点坐标(x,y) + 一维热图value值 + 二维(tag,tag')值
                    if joint[2] > 0:
                        # 该节点在对应的热力图中的坐标
                        y, x = joint[0:2]
                        xx, yy = int(x), int(y)
                        # print(batch_id, joint_id, det[batch_id].shape)
                        # 热力图det.shape=(1,17,h,w)
                        # 该节点所在的热力图tmp
                        # tmp.shape=(h,w)
                        tmp = det[batch_id][joint_id]

                        # 微调该节点的y坐标
                        if tmp[xx, min(yy+1, tmp.shape[1]-1)] > tmp[xx, max(yy-1, 0)]:
                            y += 0.25
                        else:
                            y -= 0.25

                        # 微调该节点的x坐标
                        if tmp[min(xx+1, tmp.shape[0]-1), yy] > tmp[max(0, xx-1), yy]:
                            x += 0.25
                        else:
                            x -= 0.25
                        # 更新节点坐标(x,y)
                        ans[batch_id][people_id, joint_id, 0:2] = (y+0.5, x+0.5)
        return ans

    def refine(self, det, tag, keypoints):
        """
        Given initial keypoint predictions, we identify missing joints
        :param det: numpy.ndarray of size (17, 128, 128)
        :param tag: numpy.ndarray of size (17, 128, 128) if not flip
        :param keypoints: numpy.ndarray of size (17, 4) if not flip, last dim is (x, y, det score, tag score)
        :return:
        """
        if len(tag.shape) == 3:
            # tag shape: (17, 128, 128, 1)
            tag = tag[:, :, :, None]

        tags = []
        for i in range(keypoints.shape[0]):
            if keypoints[i, 2] > 0:
                # save tag value of detected keypoint
                x, y = keypoints[i][:2].astype(np.int32)
                tags.append(tag[i, y, x])

        # mean tag of current detected people
        prev_tag = np.mean(tags, axis=0)
        ans = []

        for i in range(keypoints.shape[0]):
            # score of joints i at all position
            tmp = det[i, :, :]
            # distance of all tag values with mean tag of current detected people
            tt = (((tag[i, :, :] - prev_tag[None, None, :]) ** 2).sum(axis=2) ** 0.5)
            tmp2 = tmp - np.round(tt)

            # find maximum position
            y, x = np.unravel_index(np.argmax(tmp2), tmp.shape)
            xx = x
            yy = y
            # detection score at maximum position
            val = tmp[y, x]
            # offset by 0.5
            x += 0.5
            y += 0.5

            # add a quarter offset
            if tmp[yy, min(xx + 1, tmp.shape[1] - 1)] > tmp[yy, max(xx - 1, 0)]:
                x += 0.25
            else:
                x -= 0.25

            if tmp[min(yy + 1, tmp.shape[0] - 1), xx] > tmp[max(0, yy - 1), xx]:
                y += 0.25
            else:
                y -= 0.25

            ans.append((x, y, val))
        ans = np.array(ans)

        if ans is not None:
            for i in range(det.shape[0]):
                # add keypoint if it is not detected
                if ans[i, 2] > 0 and keypoints[i, 2] == 0:
                # if ans[i, 2] > 0.01 and keypoints[i, 2] == 0:
                    keypoints[i, :2] = ans[i, :2]
                    keypoints[i, 2] = ans[i, 2]

        return keypoints

    """
    #############################################
    # tensor张量：热力图det.shape=(1,17,h,w)
    # tensor张量：tag图tag.shape=(1,17,h,w,2)
    #############################################
    """
    def parse(self, det, tag, adjust=True, refine=True):
        """
        ####################################
        topk返回值：
        # tag_k.shape=(1,17,30,2)
        # ind_k.shape=(1,17,30,2)二维坐标(x,y)
        # val_k.shape=(1,17,30)
        ####################################
        match返回值：
        # ans为列表，图片batch_size数目len(ans)=1
        # ans=[array(num_people,17,5),] ，一张图有num_people个人 ，每个人体对应一个17*5的array数组。
        #######################################################################################
        """
        # ans为列表，图片batch_size数目len(ans)=1 ，ans=[array(num_people，17，5),]
        ans = self.match(**self.top_k(det, tag))

        if adjust:
            # 微调整每个节点的(x,y)坐标
            ans = self.adjust(ans, det)

        """
        求一个人的所有17节点的热图值的均值，
        并将每个人的热图均值加入到scores列表中。
        """
        # ans[0]是一个numpy数组，ans[0]=array(num_people,17,5)
        # len(scores)=num_people
        # scores=[ heatpoint1,heatpoint2,...,heatpoint_numpeople ]
        scores = [i[:, 2].mean() for i in ans[0]]

        if refine:
            ans = ans[0]
            # for every detected person
            for i in range(len(ans)):
                det_numpy = det[0].cpu().numpy()
                tag_numpy = tag[0].cpu().numpy()
                if not self.tag_per_joint:      #  self.tag_per_joint=True
                    tag_numpy = np.tile(
                        tag_numpy, (self.params.num_joints, 1, 1, 1)
                    )
                ans[i] = self.refine(det_numpy, tag_numpy, ans[i])
            ans = [ans]

        """
        ##################################
        ans = [ans]
        得到ans=[ array(num_people,17,5),]
        ##################################
        len(scores)=num_people
        scores=[ heatpoint1,heatpoint2,...,heatpoint_numpeople ]
        ##########################################################
        """
        return ans, scores