import numpy as np
import matplotlib.pyplot as plt
from pylab import figure,plot,matplotlib,loadtxt
import os

# Import correlation functions
from cor_funcs import *

# 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.

matplotlib.rcParams.update({'font.size': 16})
matplotlib.rc('legend',**{'fontsize': 14})

# load parameters from file
corPms = loadtxt('corParams')
ReyPwr = float(corPms[0])
intPwr = float(corPms[1])
m = float(corPms[2])
b = float(corPms[3])

heightDict = {
        '200': 101,
        '140': 80,
        '100': 57
}

# example key syntax               - '3_2_140_15_2_0'
# example name of class from xfoil - 'R1600000_0_A5X'
if __name__ == '__main__':

    #------------------- Default Figure Size ---------------------
    fig = figure(num=None, figsize=(14,6), dpi=80 )
    fig.subplots_adjust(bottom=0.12)
    fig.subplots_adjust(top=0.91)
    fig.subplots_adjust(right=0.96)
    fig.subplots_adjust(left=0.1)
    fig.subplots_adjust(wspace=0.14)
    #------------------------------------------------------------

    rtcrt,ncrit,alphs,dist,labels,keys = [],[],[],[],[],[]
    Reystr = [ '1_6','2_4','3_2' ]
    rhts   = [ '140' ]
    dens   = [ '15' ]
    # this angle range covers all of them not just ones with exp data
    aoas   = np.arange(-10,9,0.5)
    pq,crit = [],[]
    classes, exptrans = [],[]

    #m,b = -0.03313, 2.962
    #m,b = -0.04368, 3.6419
    # m,b = -0.20, 2.458
    #m,b = -0.03308, 1.8567
    for den in dens:
      subFignum = 1
      for rstr in Reystr:
      # intitialize arrays for this particular Reynolds number
        xAoAs,xTrans,xAoAsC,xTransC = [],[],[],[]
        rnum = ''.join([list(rstr)[0],list(rstr)[2]])
        for rht in rhts:
            ks = heightDict[rht]
            k_s = float(ks)*1.0e-6
            # automate by generating the substring on the fly
            # e.g. subst = '2_4_140_03_'
            subst   = ''.join([rstr, '_', rht, '_', den, '_' ])
            curkeys = get_key_list( subst,transdict )
            expAoA,expTrans,expSig = getexp(transdict,subst)
            # loop through the keys that indicate which we have data for
            for aoa in aoas:
                err    = False
                aoastr = str(aoa)
                classtr = ''.join( ['R',rnum,'00000_0_A',aoastr,'X' ] )
                cname = convert_cname( classtr )
                try: thisclass = globals()[ cname ]
                except: err = True
                if not err:
                      # extract the data from class for this configuration
                      xu,su,n,rt,Dcp,cf,K,mut,dstr,mom,thw = extract_vars( thisclass )
                      #xl,sl,n,rt,Dcp,cf,K,mut,dstr,mom,thw = extract_vars_l( thisclass )

                      xAoAsC.append(aoa)
                      xTransC.append(xu[-1])
                      # 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)**ReyPwr) )

                      # here we pull the function from the class, which has most
                      # of the information built into the function
                      foo   = getattr( thisclass, 'intquant' )
                      Ilambda = foo(0.12,0.02,np.array(quant2))**intPwr

                      # call function to pull transition location with Ilambda
                      roughTrans = pull_cor( Ilambda,m,b,thisclass.rtu,thisclass.xu )
                      xTrans.append( roughTrans )

        # generate subfigure
        subplotnum = ''.join(['13',str(subFignum)])
        ax = fig.add_subplot(subplotnum)
        subFignum += 1
        # pull the clean experimental data
        cleankey = ''.join([rstr,'_cle'])
        acln,tcln,sigc = getexp( transdict,cleankey)
        # plot the data
        ax.errorbar(tcln,acln,xerr=sigc,fmt='o',mfc='none',mew=2,color='k',capsize=4,markersize=7)
        ax.errorbar(expTrans,expAoA,xerr=expSig,fmt='o',color=[0.0,0.4,0.0], mec=[0.0,0.4,0.0],capsize=4,markersize=10)
        plt.plot(xTransC,xAoAsC,'--s',color=[0.0,0.2,0.2],mec=[0.0,0.2,0.2],mfc='none',mew=2,markersize=8,linewidth=2)
        plt.plot(xTrans,xAoAsC,'--s',color=[1.0,0.0,0.3],mec=[1.0,0.0,0.3],markersize=8,linewidth=2)
        print xTrans
        print xAoAsC
        if subFignum == 2:
            labels = ['Clean (exp)', 'Rough (exp)', 'Clean (XFOIL)', 'Rough (XFOIL)']
            plt.legend(labels,loc=3,numpoints=1)
            plt.ylabel('Angle of Attack (deg)')
        plt.xlabel('Transition Location (x/c)')
        plt.xlim([0.0,0.8])
        plt.grid()

    plt.show()
