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. 
#
#
# Primary correlation function F() -> int_(s.p)^(r.e) f() ds 
############################################################
# 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 
#    thw  -- Thwaites pressure gradient parameter (lambda_theta)
# Output:
#       quant - Some combination of input parameters f(k^+)*g(K) 
#
def comp_quant( k_s,K,mutnu,thw,dstar,mom ):
	""" Function of boundary layer quantities that will be integrated """
	quant  = [] 
	for i in range(len(K)):
    		kplus = mutnu[i]*k_s 
		# do not add if we are outside thresholds 
		if kplus < 5.0 or K[i] > 6.0e-4: 
                 quant.append( float(0.0) )
		else:
			#fthws = 1.0 / (1.0 + np.exp(100.0*(thw[i]-0.0035))) + 0.5	       		   
                 #quant.append( float( kplus**4*(1.0 + 0.4*dcp[i]) ) ) 
                 th = float( thw[i] )
                 if th <= 0:
                     ft = 1 - (-12.9*th - 123.7*th**2 - 405.7*th**3)
                 else:
                     ft = 1 + 0.275*( 1- np.exp(-35.0*th) )
                 fthws = (ft)**2
                 #fthws = K[i]
                 print fthws
                 #quant.append( (float( fthws*(np.exp(0.4*kplus))*dstar[i]**2.0) ) )  
                 quant.append( (float( (kplus)**6 )))   
	return quant



# 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()


# Function that takes in an x_class along with a target Re_theta,
# then outputs the streamwise location it occurs.
# For the time being assuming it is only the upper surface
def get_trans_ret( xclass,retcrit ):
    ret,x = xclass.rtu, xclass.xu
    rt,ct = ret[0],0
    while rt < retcrit:
        lowind  = ct
        highind = ct + 1
        ct = ct + 1
        rt = ret[ct]
    retlow,rethigh = ret[lowind],ret[highind]
    xlow,xhigh = x[lowind],x[highind]
    pct = (ret - retlow)/(rethigh - retlow)    
    xout = xlow + pct*(xhigh - xlow) 
    return xout 


# Function that takes in x_class and outputs values of interest
# Once again assume we are on the upper surface
def extract_vars( xclass ):
    xu   = np.array( getattr( xclass, 'xu'    ) )
    su   = np.array( getattr( xclass, 'su'    ) )
    n    = np.array( getattr( xclass, 'nu'    ) )
    rt   = np.array( getattr( xclass, 'rtu'   ) ) 
    Dcp  = np.array( getattr( xclass, 'Dcpu'  ) )  
    cf   = np.array( getattr( xclass, 'cfu'   ) )    
    K    = np.array( getattr( xclass, 'Ku'    ) )    
    mut  = np.array( getattr( xclass, 'mutnu' ) )    
    dstr = np.array( getattr( xclass, 'dstu'  ) )    
    mom  = np.array( getattr( xclass, 'thtu'  ) )    
    thw  = np.array( getattr( xclass, 'Twtsu' ) )    
    return xu,su,n,rt,Dcp,cf,K,mut,dstr,mom,thw 


# Function that finds the point where the slope starts to drop 
def bin_arrays( xarray,yarray ):
    span = np.amax(xarray) - np.amin(xarray)  
    bins = np.arange( np.amin(xarray), np.amax(xarray), span/50)
    data = np.zeros(len(bins)) 
    binct = np.zeros(len(bins)) 
    for i in range(len(bins)-1):
        for j in range(len(xarray)):
            if xarray[j] > bins[i] and xarray[j] < bins[i+1]:
                data[i] = data[i] + yarray[j]
                binct[i] = binct[i] + 1           
    for i in range(len(data)):
        if binct[i] > 0:
            data[i] = data[i]/binct[i] 
        else: 
            data[i] = 1.0
    binmin = [] 
    for i in range(len(data)-1):
        if data[i] < 0.9 and data[i+1] < 0.9:
            binmin.append( bins[i+1] ) 
    return np.amin( binmin )  


############################################################################
# example key syntax               - '3_2_140_15_2_0'
# example name of class from xfoil - 'R1600000_0_A5X'
#
if __name__ == '__main__':
    if os.path.exists( 'cor_data.txt' ):
    	os.remove( 'cor_data.txt' )
    
    fd = os.open('cor_data.txt', os.O_RDWR|os.O_CREAT )
    
    rtcrt,ncrit,alphs,dist,labels,keys = [],[],[],[],[],[] 
    Reystr = [ '1_6','2_4','3_2' ]
    rhts   = [ '140' ] 
    dens   = [ '15' ]
    pq,crit = [],[]
    classes, exptrans = [],[] 
    for den in dens:
      #pq,crit = [],[] 
      for rstr in Reystr:
        rnum    = ''.join([list(rstr)[0],list(rstr)[2]]) 
        for rht in rhts:
            # need to do some sort of scaling to account for lower densities
            if   den == '03' : scale = 0.9
            elif den == '06' : scale = 0.92
            elif den == '09' : scale = 0.94
            elif den == '15' : scale = 1.0
            k_s = float(rht)*1.0e-6*scale
            # automate by generating the substring on the fly 
            # subst   = '2_4_140_03_'
            # pq,crit = [],[]
            subst   = ''.join([rstr, '_', rht, '_', den, '_' ])
            curkeys = get_key_list( subst,transdict )
            for key in curkeys:
                err     = False
                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:
                      classes.append( classtr )    
                      exptrans.append( float( transdict[key][0] )  )                   
                      # extract the data from class for this configuration 
                      xu,su,n,rt,Dcp,cf,K,mut,dstr,mom,thw = extract_vars( thisclass )
            		# compute the correlation function 
                      quant = comp_quant( k_s,K,mut,thw,dstr,mom ) 
                      quant2 = [] 
                      # Add some Reynolds number scaling to the computed function 
                      for q in quant:
                          quant2.append( q*(float(rnum)**0.25) )
                          #quant2.append( q )
                      #
                      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(  1.0 - (n[-1]  -  n[ind])/n[-1] )  
                      #crit.append(  1.0 - (rt[-1]  -  rt[ind])/rt[-1] )  
                      crit.append( rt[ind]/rt[-1] )  
                      dist.append( (foo(0.12,0.02,np.array(quant2)))**0.25  )
                      pq.append(   (foo(0.12,0.02,np.array(quant2)))**0.25  )       
                      keys.append( key ) 
                      alphs.append( aoastr )
      leg    = ''.join([rnum,'_',rht])  		
      labels = labels + [leg] 		
      plt.plot( pq, crit,'*',markersize=10 )                 
    for i in range(len( ncrit )):
        wrtstr = ''.join([str(dist[i]),'   ',str(ncrit[i]),'   ',str(keys[i]),'\n'])
        os.write(fd,wrtstr)
   
    thresh = bin_arrays( pq,crit )
    print thresh
    # generate a linear best fit line
    # now we are looking at finding the line that is after the flat line 
    downslope = []
    downdist  = []
    for i in range(len(crit)):
        if dist[i] > thresh:
            downdist.append( dist[i] )
            downslope.append( crit[i] )
    
    m,c  = np.polyfit( downdist,downslope,1 )
    print m,c 
    line = []
    for d in dist:  
    	line.append( d*m + c )
    
    plt.plot( dist,line  ) 
    # plt.legend( labels )
    plt.xlabel( '$\\int \ \  R_{\\Lambda} \ ds$' )
    plt.ylabel( '$\\frac{Re_{\\theta t, rough}}{Re_{\\theta t}}  $', fontsize=22 )
    plt.ylim( [0,1.2 ] )
    #plt.xlim( [0,2000] ) 
    plt.show() 
    print np.corrcoef( downdist,downslope ) 
    

