#!/usr/bin/python

from __future__ import print_function
import sys
import os.path
import re
import math
import time
import random
import copy
import array
import numpy as np
from numpy.linalg import inv
from math import sqrt
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

sys.path.append("/mntdirect/_users/semeraro/python_tools") 
import FitTools
import FormFactor
#import Decoupling

###########################################################################################

def X2function(Q,IDATA,I,ERR,N):
	'X^2 calculation'
	X2 = np.sum(  ( ( IDATA-I ) / ERR )**2 )
	return X2/(Q.shape[0]-(N-1))
	
def exponential_decay(x, a, st, tau_iter, Cave) :
    return a * np.exp( -(x-st)/tau_iter ) + Cave	
	
###########################################################################################
###########################################################################################
def Min_likelyhood (prior_ratio) : 
	min = 1 
	if prior_ratio < 1 :
		min = prior_ratio
	return min
    
#------------------------- Partial ratio between priors
def PriorRatio (start, check, rel, prev) :
	new = np.exp(-(start-check)**2/(2*(rel*start)**2))
	old = np.exp(-(start-prev)**2/(2*(rel*start)**2))
	return new/old

###########################################################################################
###########################################################################################
###########################################################################################
###########################		SIMULATED ANNEALING ALGORITHM	###########################
###########################################################################################
###########################################################################################
###########################################################################################
#def SimAnnealing ( Q, IDATA, ERR, PAR, FIX, PRIOR, L_LIM, H_LIM, function, temperature, save ):
def SimAnnealing ( input_list ):
	'Simulated Annealing Algorithm'
	
	Q, IDATA, ERR, PAR, FIX, PRIOR, L_LIM_org, H_LIM_org, function, temperature, save = input_list

	print("START")	
	#####################################################
	#------------------------ Define Controls
	good = 0
	tau = 5 # tau>1 fast / tau<1 slow
	maxgood = 12500 
	maxcount = 50000 
	thermo = 10 # for POPC LUVs
	#thermo = 500
	T = temperature
 
	FREE_MIN = []
	
	prBol = 0 
	
	alpha		= 0.15 
	alphamax	= 1.0 
	alphamin	= 0.024 
		
	cumDEnt	= 0
	cumDX2 	= 0
	dly 	= 0 
	
	negWater_check = 0	
	
	stop = "CONTINUE"

	#####################################################
	#------------------------ Define Free Parameters Array
	(FREE,L_LIM,H_LIM) = FitTools.FreePar(PAR,FIX,L_LIM_org,H_LIM_org)
	START	= FREE.copy()

	#####################################################
	#------------------------ Merge Parameters Arrays
	CALC = FitTools.MergePar(PAR,FREE,FIX)
	
	#####################################################
	#------------------------ Evaluate I(q) & Calculate new CHI squared
	
	#I = function(Q,CALC)
	intensity_init = function(Q,CALC)
	I = intensity_init.intensity()
	X2 = X2function(Q,IDATA,I,ERR,len(FREE))
	X2_min = X2
			
	#####################################################
	#------------------------ Set priors index and sigma
	pr_idx	= []
	pr_rel	= []
	
	p = 0
	f = 0
	while p < len(PRIOR) :
		if FIX[p] == 'f' : f+= 1 				
		if PRIOR[p] != '-':
			pr_idx.append(p-f)
			pr_rel.append(float(PRIOR[p]))	
		p+=1	
	
	#print(pr_idx)
	#print(pr_rel)
		
	#####################################################
	#------------------------ Start Loop
	loop=0
	t0=time.time()	
	Nm = 500
	
	trace= [[] for i in range(3+len(PAR))]
	
	while stop == "CONTINUE":
		
		#####################################################
		#------------------------ Calculate Parameter Perturbations
		PERT=[]
		for i in range(len(FREE)):
			PERT.append( np.random.normal( 0, alpha * ( H_LIM[i] - L_LIM[i] )/6. ) )
							
		#####################################################
		#------------------------ Calculate temporary parameters
		#------------------------ Check for FIX parameters & boundarys
		CHECK=[]
		for i in range(len(FREE)):
			if (FREE[i]+PERT[i])>H_LIM[i] :
				CHECK.append(FREE[i]+(H_LIM[i]-FREE[i])/10)
			elif (FREE[i]+PERT[i])<L_LIM[i] :
				CHECK.append(FREE[i]-(FREE[i]-L_LIM[i])/10)
			else: 
				CHECK.append(FREE[i]+PERT[i])

		#####################################################
		#------------------------ Merge Temporary Parameters arrays
		CALC = FitTools.MergePar(PAR,CHECK,FIX)
		
		#####################################################
		#------------------------ Evaluate Temporary I(q) & Calculate new CHI squared
		#I_TEMP = function(Q,CALC)
		intensity_init = function(Q,CALC)
		I_TEMP = intensity_init.intensity()		

		#####################################################
		#------------------------ Check for negative water
		check_negH2O = intensity_init.negative_water()	
		X2_TEMP = X2function(Q,IDATA,I_TEMP,ERR,len(FREE))
		if check_negH2O != 0 :
			X2_TEMP*= 5.0*(1+check_negH2O)

		#####################################################
		#------------------------ Get prior ratio	
		prior = 1
		for i in range(len(pr_idx)) :
			prior*= PriorRatio( START[pr_idx[i]], CHECK[pr_idx[i]], pr_rel[i], FREE[pr_idx[i]] )
		
		#####################################################
		#------------------------ Check new X2
		#------------------------ Update configuration
		prBol = np.exp(-(X2_TEMP-X2)/T) * Min_likelyhood(prior)
		
		if X2_TEMP < X2 and negWater_check == 0 : 
			FREE = CHECK.copy()
			I = I_TEMP.copy()
			cumDX2+= ( X2_TEMP - X2 )			
			X2 = X2_TEMP
			good+= 1
			accept="Good - Accepted"
			if X2 < X2_min :
				FREE_MIN = FREE.copy()
				I_MIN = np.copy(I)
				X2_min = X2
			##### Saving the trace
			if good % 10 == 0 : 
				trace[0].append(good)
				trace[1].append(T)
				trace[2].append(X2_TEMP)
				for p in range(len(CALC)) : trace[p+3].append(CALC[p])
		elif random.random() <= prBol and negWater_check == 0 : 
			FREE = CHECK.copy()
			I = I_TEMP.copy()
			cumDX2+= ( X2_TEMP - X2 )			
			X2 = X2_TEMP
			good += 1
			accept="Bad - Accepted"
			##### Saving the trace
			if good % 10 == 0 : 
				trace[0].append(good)
				trace[1].append(T)
				trace[2].append(X2_TEMP)
				for p in range(len(CALC)) : trace[p+3].append(CALC[p])
		else :
			cumDEnt+= -( X2_TEMP - X2 ) / T	
			accept="Bad - Rejected"
	
		#####################################################
		#------------------------ Update the loop counter
		loop+= 1

		#####################################################
		#------------------------ Temperature scheme
		if cumDEnt == 0 or cumDX2 >= 0 : 
			T = temperature
		else :
			T = thermo * cumDX2 / cumDEnt
			if T < 1 : 
				T = 1
 
		if		(good+1)/(loop+1) < 0.6 : alpha-= alpha*0.1
		elif	(good+1)/(loop+1) > 0.7 : alpha+= alpha*0.1
        
		if 		alpha < alphamin : alpha = alphamin
		elif 	alpha > alphamax : alpha = alphamax
	
		#####################################################
		#------------------------ Update maxgood
		if loop % 1000 == 0 and loop < 9001 :
		
			try:
				popt, pcov = curve_fit( exponential_decay, np.array(trace[0]), np.array(trace[2]) )
				a_fit, st_fit, tau_iter_fit, Cave_fit = popt

			except RuntimeError:
				#print("\nError - curve_fit failed\n")
				tau_iter_fit = 500	

			#fig1, (figA) = plt.subplots(1, 1, figsize=(7.0, 7.0))
			#plt.scatter(trace[0],trace[2], marker='o', facecolors='none', color='b')
			#plt.plot(np.array(trace[0]),exponential_decay(np.array(trace[0]),a_fit, st_fit, tau_iter_fit, Cave_fit), '-', color='r', linewidth=2)
			#fig1.tight_layout()
			#plt.xscale('log')
			#plt.yscale('log')			
			#plt.show()
	
			if		tau_iter_fit > 500 :	maxgood = 12500
			elif	tau_iter_fit*5 < 250 :	maxgood = 10250			
			else :							maxgood = int(tau_iter_fit*5)  + 10000
			
		#####################################################
		#------------------------ Time counting
		t1=time.time()-t0
		(minutes,seconds,cents)=FitTools.TimeCount(t1)
		#success = (good+1)/(loop+1)
		
		#if		loop/1000 == 1: dly = (time.time()-t0)/1001
		#elif	loop/3000 == 1: dly = (time.time()-t0)/3001
		#elif	loop/6000 == 1: dly = (time.time()-t0)/6001
		#elif	loop/9000 == 1: dly = (time.time()-t0)/9001	
		#(minutes_ex,seconds_ex,cents_ex)=FitTools.TimeCount( dly * maxgood/success  )#*maxgood/((good+1)/(loop+1))-t0)					

		#####################################################
		#------------------------ Print Loop Results	
		#if T>0.9 and loop%10 == 0 :
		#	out="T = %0.3f \ a = %0.3f \ min.X²= %0.5s \ N. %0.3d/%0.3d (%0.3f) -> %.3d \ %2d:%2d s (exp. %2d:%2d s) \ " %(T,alpha,X2_min,good,loop,success,maxgood,minutes,seconds,minutes_ex,seconds_ex)
			#print (out, end='\r')	
			
		#####################################################
		#------------------------ Set the Loop Conditions
		if good == maxgood+1 or loop == maxcount+1 : 
			stop="STOP"
			print("STOP")
	
	#####################################################
	#------------------------ End Loop
	#with open("./"+save+"/Trace.dat", 'w') as trc:
	#	for i in range(len(trace[p])) :
	#		for p in range(len(trace)) :
		#		trc.writelines( "%.7f\t" %(trace[p][i]) ) 	
	#		trc.writelines( "\n" ) 		


	return [ FitTools.MergePar(PAR,FREE_MIN,FIX), X2_min]
	

###########################################################################################
###########################################################################################