import numpy as np
# can do something fancier here to automate string generation 
import matplotlib.pyplot as plt
import os

# Import all the data 
from transdict_err import *
from NXdata        import * 


# Now we have a dictionary that has all the transition locations and a collection of classes that have all the BLinfo generated by XFOIL. 


# Function that takes a string and a dictionary, then outputs a list of all the keys in the dictionary that contain that string
def get_key_list( substring,inpdict ):
	outkeys = [];
	for key in inpdict.keys():
		if substring in key: outkeys.append( key )
	return outkeys


# take a string and a substring a return the difference 
def get_str_diff( string,substring ):
	return list(string)[len(list(substring)):] 
	

# take in array and scalar, then return the index of the array 
# that is nearest in value to scalar 
def near_ind( array,value ):
	return (np.abs(array-value)).argmin()


############################################################
# Inputs:
#	k_s  -- Some "equivalent" sand grain roughness height
#	K    -- Relaminarization parameter (nu/U_e)*(dU_e/ds)
#       mutnu - Frictional velocity divided by nu (k^+ = mutnu*k_s) 
#	dstar - Boundary layer diplacement thickness 
# Output:
#       quant - Some combination of input parameters f(k^+)*g(K) 
#
def comp_quant( k_s,K,mutnu,thw,dst,sig,EX ):
	""" Money Function """
	quant  = [] 
	#C1     =  -100.0
	#C2     =  0.0
	for i in range(len(K)):
		kplus = mutnu[i]*k_s 
		# do not add if we are outside thresholds 
		if kplus < 1.0 or K[i] > 6.0e-3:
			quant.append( float(0.0) )
		else:
			#quant.append( float( kplus**4*(1.0 + 0.4*dcp[i]) ) )
                 # create some function of the Thwaites parameter
                 fthws = 1.0 / (1.0 + np.exp(sig[0]*(thw[i]-sig[1]))) + sig[2]
                 quant.append( (fthws*kplus**EX)*dst[i]**2 )  
	return quant

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

# example key syntax               - '3_2_140_15_2_0'
# example name of class from xfoil - 'R1600000_0_A5X'

if os.path.exists( 'cor_sweep.txt' ):
	os.remove( 'cor_sweep.txt' )

fd = os.open('cor_sweep.txt', os.O_RDWR|os.O_CREAT )
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#           !!! MASSIVE PARAMETER SWEEP !!! 
# Think of what values the variables should take...

Sig1a = np.arange( 0,50,5 )
Sig2a = np.arange( 0.0, 0.005, 0.0005) #Should really make this go to zero
Sig3a = np.arange( 0.5, 0.6, 0.1 )
sig   = []

for i in range( len(Sig1a) ):
    for j in range( len(Sig2a) ):
        for k in range( len(Sig3a) ):
          sig.append( [Sig1a[i],Sig2a[j],Sig3a[k]] ) 

Expa = np.arange(4,6.3,1.0) 
Rta  = np.arange(0.4,0.6,0.2) 

# create a multi dimensional array with all the parameters we are sweeping
params = []
for i in range( len(sig) ): 
	for j in range( len(Expa) ):
		for k in range( len(Rta) ):
			params.append([ sig[i],Expa[j],Rta[k] ]) 


numconfigs = len( params ) 
# Test for now
# numconfigs = 400
print numconfigs

for i in range( numconfigs ):
    print i
    rtcrt,ncrit,alphs,dist,labels,keys = [],[],[],[],[],[] 
    Reystr = [ '1_6','2_4','3_2' ]
    rhts   = [ '100','140' ] 
    for rstr in Reystr:
        rnum    = ''.join([list(rstr)[0],list(rstr)[2]]) 
        for rht in rhts:
    	    k_s = float(rht)*1.0e-6
            pq,crit =[],[] 
            # automate by generating the substring on the fly 
            # subst   = '2_4_140_03_'
    	    subst   = ''.join([rstr,'_',rht,'_15_' ])
            curkeys = get_key_list( subst,transdict )
            for key in curkeys:
            	err     = False
    		#aoastr = '2_0' 
            	aoastr  = ''.join( get_str_diff( key,subst ))
            	classtr = ''.join( ['R',rnum,'00000_0_A',aoastr,'X' ] )
            	try: thisclass = globals()[ classtr ]
            	except: err = True
            	if not err:
            	 	# extract the data from class for this configuration 
            		xu   = np.array( getattr( thisclass, 'xu'    ) )
            		su   = np.array( getattr( thisclass, 'su'    ) )
            		n    = np.array( getattr( thisclass, 'nu'    ) )
            		rt   = np.array( getattr( thisclass, 'rtu'   ) )
            		Dcp  = np.array( getattr( thisclass, 'Dcpu'  ) )  
            		cf   = np.array( getattr( thisclass, 'cfu'   ) )    
            		K    = np.array( getattr( thisclass, 'Ku'    ) )    
            		mut  = np.array( getattr( thisclass, 'mutnu' ) )    
            		dstr = np.array( getattr( thisclass, 'dstu'  ) )    
            		thw  = np.array( getattr( thisclass, 'Twtsu' ) )    
            		# compute the correlation function 
			sig  = params[i][0]
			EX   = params[i][1] 
    	               	quant = comp_quant( k_s,K,mut,thw,dstr,sig,EX ) 
    			quant2 = [] 
    			for q in quant:
    				quant2.append( q*(float(rnum)/10.0) )
            		foo   = getattr( thisclass, 'intquant' ) 
            		trans = float( transdict[key][0] ) 
            		ind   = near_ind( xu,trans )
            		ncrit.append( 1.0 - (n[-1]  -  n[ind])/n[-1] )
            		rtcrt.append( 1.0 - (rt[-1] - rt[ind])/rt[-1] ) 
    			crit.append(  n[ind]  )  
    			#crit.append(  1.0 - (n[-1]  -  n[ind])/n[-1] )  
    			#crit.append(  1.0 - (rt[-1]  -  rt[ind])/rt[-1] )
			root = params[i][2] 
            		dist.append( (foo(0.12,0.02,np.array(quant2)))**root  )
            		pq.append(   (foo(0.12,0.02,np.array(quant2)))**root  )
    			keys.append( key ) 
            		alphs.append( aoastr )
    corcoef  = np.corrcoef( ncrit, dist ) 
    writestr = ''.join( [ str(params[i]),'          ',str(corcoef[0][1]),'\n' ] ) 
    os.write(fd,writestr)



















