import nibabel
import os
import numpy as np
import dipy as dp
import dipy.core.generalized_q_sampling as dgqs
import dipy.io.pickles as pkl
import scipy as sp
from matplotlib.mlab import find
import dipy.core.sphere_plots as splots
import dipy.core.sphere_stats as sphats
import dipy.core.geometry as geometry
import get_vertices as gv

#old SimData files
'''
results_SNR030_1fibre
results_SNR030_1fibre+iso
results_SNR030_2fibres_15deg
results_SNR030_2fibres_30deg
results_SNR030_2fibres_60deg
results_SNR030_2fibres_90deg
results_SNR030_2fibres+iso_15deg
results_SNR030_2fibres+iso_30deg
results_SNR030_2fibres+iso_60deg
results_SNR030_2fibres+iso_90deg
results_SNR030_isotropic
'''
#fname='/home/ian/Data/SimData/results_SNR030_1fibre'
''' file  has one row for every voxel, every voxel is repeating 1000
times with the same noise level , then we have 100 different
directions. 1000 * 100 is the number of all rows.

The 100 conditions are given by 10 polar angles (in degrees) 0, 20, 40, 60, 80,
80, 60, 40, 20 and 0, and each of these with longitude angle 0, 40, 80,
120, 160, 200, 240, 280, 320, 360.
'''

#new complete SimVoxels files
simdata = ['fibres_2_SNR_80_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_60_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_40_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_40_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_20_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_100_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_20_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_40_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_60_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_100_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_60_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_80_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_100_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_100_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_80_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_60_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_40_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_80_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_20_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_60_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_100_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_100_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_20_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_1_SNR_20_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_40_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_20_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_80_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_80_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_20_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_60_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_100_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_80_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_60_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_20_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_100_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_20_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_80_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_80_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_100_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_40_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_1_SNR_60_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_40_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_60_angle_30_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_40_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_60_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_80_angle_15_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_1_SNR_40_angle_00_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_100_angle_60_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00',
 'fibres_2_SNR_40_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_1_diso_0.7',
 'fibres_2_SNR_20_angle_90_l1_1.4_l2_0.35_l3_0.35_iso_0_diso_00']

simdir = '/home/ian/Data/SimVoxels/'

def gq_tn_calc_save():

    for simfile in simdata:

        dataname = simfile
        print dataname

        sim_data=np.loadtxt(simdir+dataname)

        marta_table_fname='/home/ian/Data/SimData/Dir_and_bvals_DSI_marta.txt'
        b_vals_dirs=np.loadtxt(marta_table_fname)
        bvals=b_vals_dirs[:,0]*1000
        gradients=b_vals_dirs[:,1:]

        gq = dp.GeneralizedQSampling(sim_data,bvals,gradients)
        gqfile = simdir+'gq/'+dataname+'.pkl'
        pkl.save_pickle(gqfile,gq)

        '''
        gq.IN               gq.__doc__          gq.glob_norm_param
        gq.QA               gq.__init__         gq.odf
        gq.__class__        gq.__module__       gq.q2odf_params
        '''

        tn = dp.Tensor(sim_data,bvals,gradients)
        tnfile = simdir+'tn/'+dataname+'.pkl'
        pkl.save_pickle(tnfile,tn)

        '''
        tn.ADC               tn.__init__          tn._getevals
        tn.B                 tn.__module__        tn._getevecs
        tn.D                 tn.__new__           tn._getndim
        tn.FA                tn.__reduce__        tn._getshape
        tn.IN                tn.__reduce_ex__     tn._setevals
        tn.MD                tn.__repr__          tn._setevecs
        tn.__class__         tn.__setattr__       tn.adc
        tn.__delattr__       tn.__sizeof__        tn.evals
        tn.__dict__          tn.__str__           tn.evecs
        tn.__doc__           tn.__subclasshook__  tn.fa
        tn.__format__        tn.__weakref__       tn.md
        tn.__getattribute__  tn._evals            tn.ndim
        tn.__getitem__       tn._evecs            tn.shape
        tn.__hash__          tn._getD
        '''

        ''' file  has one row for every voxel, every voxel is repeating 1000
        times with the same noise level , then we have 100 different
        directions. 100 * 1000 is the number of all rows.

        At the moment this module is hardwired to the use of the EDS362
        spherical mesh. I am assumung (needs testing) that directions 181 to 361
        are the antipodal partners of directions 0 to 180. So when counting the
        number of different vertices that occur as maximal directions we wll map
        the indices modulo 181.
        '''

def analyze_maxima(indices, max_dirs, subsets):
    '''This calculates the eigenstats for each of the replicated batches
    of the simulation data
    '''

    results = []

    for direction in subsets:

        batch = max_dirs[direction,:,:]

        index_variety = np.array([len(set(np.remainder(indices[direction,:],181)))])

        #normed_centroid, polar_centroid, centre, b1 = sphats.eigenstats(batch)
        centre, b1 = sphats.eigenstats(batch)

        # make azimuth be in range (0,360) rather than (-180,180)
        centre[1] += 360*(centre[1] < 0)

        #results.append(np.concatenate((normed_centroid, polar_centroid, centre, b1, index_variety)))
        results.append(np.concatenate((centre, b1, index_variety)))

    return results

#dt_first_directions = tn.evecs[:,:,0].reshape((100,1000,3))
# these are the principal directions for the full set of simulations

#gq_tn_calc_save()

eds=np.load(os.path.join(os.path.dirname(dp.__file__),'core','matrices','evenly_distributed_sphere_362.npz'))

odf_vertices=eds['vertices']

def run_comparisons(sample_data=35):
    for simfile in [simdata[sample_data]]:

        dataname = simfile
        print dataname

        sim_data=np.loadtxt(simdir+dataname)

        gqfile = simdir+'gq/'+dataname+'.pkl'
        gq =  pkl.load_pickle(gqfile)
        tnfile = simdir+'tn/'+dataname+'.pkl'
        tn =  pkl.load_pickle(tnfile)

        dt_first_directions_in=odf_vertices[tn.IN]

        dt_indices = tn.IN.reshape((100,1000))
        dt_results = analyze_maxima(dt_indices, dt_first_directions_in.reshape((100,1000,3)),range(10,90))

        gq_indices = np.array(gq.IN[:,0],dtype='int').reshape((100,1000))

        gq_first_directions_in=odf_vertices[np.array(gq.IN[:,0],dtype='int')]

        #print gq_first_directions_in.shape

        gq_results = analyze_maxima(gq_indices, gq_first_directions_in.reshape((100,1000,3)),range(10,90))

        #for gqi see example dicoms_2_tracks gq.IN[:,0]

        np.set_printoptions(precision=3, suppress=True, linewidth=200, threshold=5000)

        out = open('/home/ian/Data/SimVoxels/Out/'+'***_'+dataname,'w')

        #print np.vstack(dt_results).shape, np.vstack(gq_results).shape

        results = np.hstack((np.vstack(dt_results), np.vstack(gq_results)))
        #print results.shape
        #results = np.vstack(dt_results)

        print >> out, results[:,:]

        out.close()

        #up = dt_batch[:,2]>= 0

        #splots.plot_sphere(dt_batch[up], 'batch '+str(direction))

        #splots.plot_lambert(dt_batch[up],'batch '+str(direction), centre)

        #spread = gq.q2odf_params e,v = np.linalg.eigh(np.dot(spread,spread.transpose())) effective_dimension = len(find(np.cumsum(e) > 0.05*np.sum(e))) #95%

        #rotated = np.dot(dt_batch,evecs)

        #rot_evals, rot_evecs =  np.linalg.eig(np.dot(rotated.T,rotated)/rotated.shape[0])

        #eval_order = np.argsort(rot_evals)

        #rotated = rotated[:,eval_order]

        #up = rotated[:,2]>= 0

        #splot.plot_sphere(rotated[up],'first1000')

        #splot.plot_lambert(rotated[up],'batch '+str(direction))

def run_gq_sims(sample_data=[35,23,46,39,40,10,37,27,21,20]):

    results = []

    out = open('/home/ian/Data/SimVoxels/Out/'+'npa+fa','w')

    for j in range(len(sample_data)):

        sample = sample_data[j]

        simfile = simdata[sample]

        dataname = simfile
        print dataname

        sim_data=np.loadtxt(simdir+dataname)

        marta_table_fname='/home/ian/Data/SimData/Dir_and_bvals_DSI_marta.txt'
        b_vals_dirs=np.loadtxt(marta_table_fname)
        bvals=b_vals_dirs[:,0]*1000
        gradients=b_vals_dirs[:,1:]

        for j in np.vstack((np.arange(100)*1000,np.arange(100)*1000+1)).T.ravel():
        # 0,1,1000,1001,2000,2001,...

            s = sim_data[j,:]

            gqs = dp.GeneralizedQSampling(s.reshape((1,102)),bvals,gradients,Lambda=3.5)
            tn = dp.Tensor(s.reshape((1,102)),bvals,gradients,fit_method='LS')

            t0, t1, t2, npa = gqs.npa(s, width = 5)

            print >> out, dataname, j, npa, tn.fa()[0]

            '''
            for (i,o) in enumerate(gqs.odf(s)):
                print i,o

            for (i,o) in enumerate(gqs.odf_vertices):
                print i,o
            '''
            #o = gqs.odf(s)
            #v = gqs.odf_vertices
            #pole = v[t0[0]]
            #eqv = dgqs.equatorial_zone_vertices(v, pole, 5)
            #print 'Number of equatorial vertices: ', len(eqv)
            #print np.max(o[eqv]),np.min(o[eqv])
            #cos_e_pole = [np.dot(pole.T, v[i]) for i in eqv]
            #print np.min(cos1), np.max(cos1)
            #print 'equatorial max in equatorial vertices:', t1[0] in eqv
            #x =  np.cross(v[t0[0]],v[t1[0]])
            #x = x/np.sqrt(np.sum(x**2))
            #print x
            #ptchv = dgqs.patch_vertices(v, x, 5)
            #print len(ptchv)
            #eqp = eqv[np.argmin([np.abs(np.dot(v[t1[0]].T,v[p])) for p in eqv])]
            #print (eqp, o[eqp])
            #print t2[0] in ptchv, t2[0] in eqv
            #print np.dot(pole.T, v[t1[0]]), np.dot(pole.T, v[t2[0]])
            #print ptchv[np.argmin([o[v] for v in ptchv])]

            #gq_indices = np.array(gq.IN[:,0],dtype='int').reshape((100,1000))

            #gq_first_directions_in=odf_vertices[np.array(gq.IN[:,0],dtype='int')]

            #print gq_first_directions_in.shape

            #gq_results = analyze_maxima(gq_indices, gq_first_directions_in.reshape((100,1000,3)),range(100))

            #for gqi see example dicoms_2_tracks gq.IN[:,0]

            #np.set_printoptions(precision=6, suppress=True, linewidth=200, threshold=5000)

            #out = open('/home/ian/Data/SimVoxels/Out/'+'+++_'+dataname,'w')

            #results = np.hstack((np.vstack(dt_results), np.vstack(gq_results)))
            #results = np.vstack(dt_results)

            #print >> out, results[:,:]

    out.close()

run_comparisons()
#run_gq_sims()


"""
================================
Recognizing hand-written digits
================================

An example showing how the scikit-learn can be used to recognize images of
hand-written digits.

This example is commented in the
:ref:`tutorial section of the user manual <introduction>`.

"""
print(__doc__)

# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
# License: BSD 3 clause

# Standard scientific Python imports
import matplotlib.pyplot as plt

# Import datasets, classifiers and performance metrics
from sklearn import datasets, svm, metrics

# The digits dataset
digits = datasets.load_digits()

# The data that we are interested in is made of 8x8 images of digits, let's
# have a look at the first 3 images, stored in the `images` attribute of the
# dataset.  If we were working from image files, we could load them using
# pylab.imread.  Note that each image must have the same size. For these
# images, we know which digit they represent: it is given in the 'target' of
# the dataset.
images_and_labels = list(zip(digits.images, digits.target))
for index, (image, label) in enumerate(images_and_labels[:4]):
    plt.subplot(2, 4, index + 1)
    plt.axis('off')
    plt.imshow(image, cmap=plt.cm.gray_r, interpolation='nearest')
    plt.title('Training: %i' % label)

# To apply a classifier on this data, we need to flatten the image, to
# turn the data in a (samples, feature) matrix:
n_samples = len(digits.images)
data = digits.images.reshape((n_samples, -1))

# Create a classifier: a support vector classifier
classifier = svm.SVC(gamma=0.001)

# We learn the digits on the first half of the digits
classifier.fit(data[:n_samples / 2], digits.target[:n_samples / 2])

# Now predict the value of the digit on the second half:
expected = digits.target[n_samples / 2:]
predicted = classifier.predict(data[n_samples / 2:])

print("Classification report for classifier %s:\n%s\n"
      % (classifier, metrics.classification_report(expected, predicted)))
print("Confusion matrix:\n%s" % metrics.confusion_matrix(expected, predicted))

images_and_predictions = list(zip(digits.images[n_samples / 2:], predicted))
for index, (image, prediction) in enumerate(images_and_predictions[:4]):
    plt.subplot(2, 4, index + 5)
    plt.axis('off')
    plt.imshow(image, cmap=plt.cm.gray_r, interpolation='nearest')
    plt.title('Prediction: %i' % prediction)

plt.show()


#! /usr/bin/python3
import pylab as plb
import numpy as np

from matplotlib import rc
rc('text', usetex=True)
rc('text.latex', unicode=True)
rc('text.latex', preamble=r'\usepackage[russian]{babel}')
#rc('font',**{'family':'serif'})
rc('font',**{'size':'19'})
res = np.loadtxt('result.txt', delimiter=', ')[0:7]
#import pdb; pdb.set_trace()
#plb.barh(y_pos, performance, xerr=error, align='center', alpha=0.4)
#plb.yscale('linear')
plb.errorbar(res[:, 1], res[:, 5], yerr=res[:, 6], label='Традиционная вейвлет-сеть', linestyle='--', marker='*', color='black')
plb.errorbar(res[:, 1], res[:, 11], yerr=res[:, 12], label='Полиморфная вейвлет-сеть', marker='o', color='green')
plb.errorbar(res[:, 1], res[:, 1],  yerr=res[:, 2], label='Отношение сигнал/шум для временного ряда $d(t), S$', color='blue')
#import pdb; pdb.set_trace()
plb.fill_between(res[:, 1], res[:, 1], res[:, 1]-np.max(res[:, 1]), res[:, 1], alpha=0.1, color='blue')
plb.xscale('log')
plb.legend(loc=0)
plb.xlim(res[-1, 1]-0.1, res[0, 1]+20)
plb.ylim(0, 670)
plb.gca().set_xticks(res[:, 1])
#plb.gca().xaxis.set_major_locator(plb.LogLocator(numticks=50))
plb.gca().xaxis.set_major_formatter(plb.ScalarFormatter())
plb.ylabel('Отношение сигнал/шум для временного ряда $\hat{y}(t), M$')
plb.xlabel('Отношение сигнал/шум для временного ряда $d(t), S$')
plb.annotate('Область применения вейвлет-сетей', [7, 310])
plb.show()
polym_higest=res[:, 11]>res[:, 1]
polym_avg=res[polym_higest, 11][1:-2]
std_higest=res[:, 5]>res[:, 1]
std_avg=res[std_higest, 5][:-2]
inp_avg=res[std_higest, 1][:-2]
polym_min=res[polym_higest, 11][1:-2]-res[polym_higest, 12][1:-2]
polym_max=res[polym_higest, 11][1:-2]+res[polym_higest, 12][1:-2]
std_min=res[std_higest, 5][:-2]-res[std_higest, 6][:-2]
std_max=res[std_higest, 5][:-2]+res[std_higest, 6][:-2]
print('Улучшение в среднем на {}%'.format(np.average((polym_avg-std_avg)/std_avg*100)))
print('Улучшение в по диапазону на {0}-{1}%'.format(np.average((polym_min-std_min)/std_min*100),
      np.average((polym_max-std_max)/std_max*100)))
polym_avg_db=10*np.log10(polym_avg-inp_avg)
std_avg_db=10*np.log10(std_avg-inp_avg)
print('Улучшение в среднем на {}дб'.format(np.average(polym_avg_db-std_avg_db)))


"""

"""

import os
import sys
import argparse

import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as plt

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('-f', help='Genome average output file (from genera_per_phage_protein.py', default='/home/redwards/Desktop/gav_all_host.out')
    parser.add_argument('-n', help='taxonomy name one of: kingdom / phylum / genus / species', default='genus')
    parser.add_argument('-v', help='verbose output', action="store_true")

    args = parser.parse_args()

    ynames = {'kingdom' : 'kingdoms', 'phylum' : 'phyla', 'genus' : 'genera', 'species' : 'species'}

    col = None
    colkey = {'kingdom' : 3, 'phylum' : 4, 'genus' : 5, 'species' : 6}
    if args.n not in colkey:
        sys.stderr.write("Sorry, taxonomy name must be one of {}\n".format("|".join(list(colkey.keys()))))
        sys.exit(-1)
    col = colkey[args.n]

    want = {'Gut', 'Mouth', 'Nose', 'Skin', 'Lungs'}

    data = {}
    with open(args.f, 'r') as fin:
        for l in fin:
            p=l.strip().split("\t")
            if p[2] not in want:
                p[2] = 'All phages'
                #continue  ## comment or uncomment this to include/exclude all data
            if p[2] not in data:
                data[p[2]] = []
            data[p[2]].append(float(p[col]))

    labels = sorted(data.keys())
    scores = []
    count = 1
    ticks = []
    for l in labels:
        scores.append(data[l])
        ticks.append(count)
        count += 1

    fig = plt.figure()
    ax = fig.add_subplot(111)

    # ax.boxplot(alldata)
    vp = ax.violinplot(scores, showmeans=True)
    for i, j in enumerate(vp['bodies']):
        if i == 0:
            j.set_color('gray')
        elif i == 1:
            j.set_color('sandybrown')
        else:
            j.set_color('lightpink')

    ax.set_xlabel("Body Site")
    ax.set_ylabel("Average number of {}".format(ynames[args.n]))
    ax.set_xticks(ticks)
    ax.set_xticklabels(labels, rotation='vertical')
    ax.get_xaxis().tick_bottom()
    ax.get_yaxis().tick_left()
    fig.set_facecolor('white')

    plt.tight_layout()
    #plt.show()
    fig.savefig("/home/redwards/Desktop/bodysites.png")


#!/bin/python
'''
fluxCal.py

Created by Seth Meeker on 11-21-2012
Modified on 02-16-2015 to perform absolute fluxCal with point sources

Opens ARCONS observation of a spectrophotometric standard star and
associated wavelength cal file, reads in all photons and converts to energies.
Bins photons to generate a spectrum, then divides this into the known spectrum
of the object to create a Sensitivity curve.  This curve is then written out to
h5 file.

Flags are associated with each pixel - see headers/pipelineFlags
for descriptions. Note some flags are set here, others are set
later on when creating photon lists.
'''

import sys,os
import tables
import numpy as np
from scipy import interpolate
from scipy.optimize.minpack import curve_fit
import matplotlib.pyplot as plt
from photometry import LightCurve
from util.FileName import FileName
from util.ObsFile import ObsFile
from util import MKIDStd
from util.readDict import readDict
from util.utils import rebin
from util.utils import gaussianConvolution
from util.utils import makeMovie
from util.utils import fitBlackbody
import hotpix.hotPixels as hp
from scipy.optimize.minpack import curve_fit
from scipy import interpolate
import matplotlib
from matplotlib.backends.backend_pdf import PdfPages
from headers import pipelineFlags
import figureHeader

class FluxCal:
    def __init__(self,paramFile,plots=False,verbose=False):
        """
        Opens flux file, prepares standard spectrum, and calculates flux factors for the file.
        Method is provided in param file. If 'relative' is selected, an obs file with standard star defocused over
        the entire array is expected, with accompanying sky file to do sky subtraction.
        If any other method is provided, 'absolute' will be done by default, wherein a point source is assumed
        to be present. The obs file is then broken into spectral frames with photometry (psf or aper) performed
        on each frame to generate the ARCONS observed spectrum.
        """
        self.verbose=verbose
        self.plots = plots

        self.params = readDict()
        self.params.read_from_file(paramFile)

        run = self.params['run']
        sunsetDate = self.params['fluxSunsetLocalDate']
        self.fluxTstamp = self.params['fluxTimestamp']
        skyTstamp = self.params['skyTimestamp']
        wvlSunsetDate = self.params['wvlCalSunsetLocalDate']
        wvlTimestamp = self.params['wvlCalTimestamp']
        flatCalFileName = self.params['flatCalFileName']
        needTimeAdjust = self.params['needTimeAdjust']
        self.deadtime = float(self.params['deadtime']) #from firmware pulse detection
        self.timeSpacingCut = self.params['timeSpacingCut']
        bLoadBeammap = self.params.get('bLoadBeammap',False)
        self.method = self.params['method']
        self.objectName = self.params['object']
        self.r = float(self.params['energyResolution'])
        self.photometry = self.params['photometry']
        self.centroidRow = self.params['centroidRow']
        self.centroidCol = self.params['centroidCol']
        self.aperture = self.params['apertureRad']
        self.annulusInner = self.params['annulusInner']
        self.annulusOuter = self.params['annulusOuter']
        self.collectingArea = self.params['collectingArea']
        self.startTime = self.params['startTime']
        self.intTime = self.params['integrationTime']

        fluxFN = FileName(run=run,date=sunsetDate,tstamp=self.fluxTstamp)
        self.fluxFileName = fluxFN.obs()
        self.fluxFile = ObsFile(self.fluxFileName)

        if self.plots:
            self.plotSavePath = os.environ['MKID_PROC_PATH']+os.sep+'fluxCalSolnFiles'+os.sep+run+os.sep+sunsetDate+os.sep+'plots'+os.sep
            if not os.path.exists(self.plotSavePath): os.mkdir(self.plotSavePath)
            if self.verbose: print "Created directory %s"%self.plotSavePath

        obsFNs = [fluxFN]
        self.obsList = [self.fluxFile]

        if self.startTime in ['',None]: self.startTime=0
        if self.intTime in ['',None]: self.intTime=-1

        if self.method=="relative":
            try:
                print "performing Relative Flux Calibration"
                skyFN = FileName(run=run,date=sunsetDate,tstamp=skyTstamp)
                self.skyFileName = skyFN.obs()
                self.skyFile = ObsFile(self.skyFileName)
                obsFNs.append(skyFN)
                self.obsList.append(self.skyFile)
            except:
                print "For relative flux calibration a sky file must be provided in param file"
                self.__del__()
        else:
            self.method='absolute'
            print "performing Absolute Flux Calibration"

        if self.photometry not in ['aperture','PSF']: self.photometry='PSF' #default to PSF fitting if no valid photometry selected

        timeMaskFileNames = [fn.timeMask() for fn in obsFNs]
        timeAdjustFileName = FileName(run=run).timeAdjustments()

        #make filename for output fluxCalSoln file
        self.fluxCalFileName = FileName(run=run,date=sunsetDate,tstamp=self.fluxTstamp).fluxSoln()
        print "Creating flux cal: %s"%self.fluxCalFileName

        if wvlSunsetDate != '':
            wvlCalFileName = FileName(run=run,date=wvlSunsetDate,tstamp=wvlTimestamp).calSoln()
        if flatCalFileName =='':
            flatCalFileName=FileName(obsFile=self.fluxFile).flatSoln()

        #load cal files for flux file and, if necessary, sky file
        for iObs,obs in enumerate(self.obsList):
            if bLoadBeammap:
                print 'loading beammap',os.environ['MKID_BEAMMAP_PATH']
                obs.loadBeammapFile(os.environ['MKID_BEAMMAP_PATH'])
            if wvlSunsetDate != '':
                obs.loadWvlCalFile(wvlCalFileName)
            else:
                obs.loadBestWvlCalFile()

            obs.loadFlatCalFile(flatCalFileName)
            obs.setWvlCutoffs(-1,-1)

            if needTimeAdjust:
                obs.loadTimeAdjustmentFile(timeAdjustFileName)
            timeMaskFileName = timeMaskFileNames[iObs]
            print timeMaskFileName

            if not os.path.exists(timeMaskFileName):
                print 'Running hotpix for ',obs
                hp.findHotPixels(obsFile=obs,outputFileName=timeMaskFileName,fwhm=np.inf,useLocalStdDev=True)
                print "Flux cal/sky file pixel mask saved to %s"%(timeMaskFileName)
            obs.loadHotPixCalFile(timeMaskFileName)
            if self.verbose: print "Loaded hot pixel file %s"%timeMaskFileName

        #get flat cal binning information since flux cal will need to match it
        self.wvlBinEdges = self.fluxFile.flatCalFile.root.flatcal.wavelengthBins.read()
        self.nWvlBins = self.fluxFile.flatWeights.shape[2]
        self.binWidths = np.empty((self.nWvlBins),dtype=float)
        self.binCenters = np.empty((self.nWvlBins),dtype=float)
        for i in xrange(self.nWvlBins):
            self.binWidths[i] = self.wvlBinEdges[i+1]-self.wvlBinEdges[i]
            self.binCenters[i] = (self.wvlBinEdges[i]+(self.binWidths[i]/2.0))

        if self.method=='relative':
            print "Extracting ARCONS flux and sky spectra"
            self.loadRelativeSpectrum()
            print "Flux Spectrum loaded"
            self.loadSkySpectrum()
            print "Sky Spectrum loaded"
        elif self.method=='absolute':
            print "Extracting ARCONS point source spectrum"
            self.loadAbsoluteSpectrum()

        print "Loading standard spectrum"
        try:
            self.loadStdSpectrum(self.objectName)
        except KeyError:
            print "Invalid spectrum object name"
            self.__del__()
            sys.exit()

        print "Generating sensitivity curve"
        self.calculateFactors()
        print "Sensitivity Curve calculated"
        print "Writing fluxCal to file %s"%self.fluxCalFileName
        self.writeFactors(self.fluxCalFileName)

        if self.plots: self.makePlots()

        print "Done"

    def __del__(self):
        try:
            self.fluxFile.close()
            self.calFile.close()
        except AttributeError:#fluxFile was never defined
            pass

    def getDeadTimeCorrection(self, obs): #WRONG RIGHT NOW. NEEDS TO HAVE RAW COUNTS SUMMED, NOT CUBE WHICH EXCLUDES NOISE TAIL
        if self.verbose: print "Making raw cube to get dead time correction"
        cubeDict = obs.getSpectralCube(firstSec=self.startTime, integrationTime=self.intTime, weighted=False, fluxWeighted=False)

        cube= np.array(cubeDict['cube'], dtype=np.double)
        wvlBinEdges= cubeDict['wvlBinEdges']
        effIntTime= cubeDict['effIntTime']
        if self.verbose: print "median effective integration time = ", np.median(effIntTime)

        nWvlBins=len(wvlBinEdges)-1
        if self.verbose: print "cube shape ", np.shape(cube)
        if self.verbose: print "effIntTime shape ", np.shape(effIntTime)

        #add third dimension to effIntTime for  broadcasting
        effIntTime = np.reshape(effIntTime,np.shape(effIntTime)+(1,))
        #put cube into counts/s in each pixel
        cube /= effIntTime

        #CALCULATE DEADTIME CORRECTION
        #NEED TOTAL COUNTS PER SECOND FOR EACH PIXEL TO DO PROPERLY
        #ASSUMES SAME CORRECTION FACTOR APPLIED FOR EACH WAVELENGTH, MEANING NO WL DEPENDANCE ON DEAD TIME EFFECT
        DTCorr = np.zeros((np.shape(cube)[0],np.shape(cube)[1]),dtype=float)
        for f in range(0,np.shape(cube)[2]):
            #if self.verbose: print cube[:,:,f]
            #if self.verbose: print '-----------------------'
            DTCorr += cube[:,:,f]
            #if self.verbose: print DTCorr
            #if self.verbose: print '\n=====================\n'
        #Correct for firmware dead time (100us in 2012 ARCONS firmware)
        DTCorrNew=DTCorr/(1-DTCorr*self.deadtime)
        CorrFactors = DTCorrNew/DTCorr #This is what the frames need to be multiplied by to get their true values
        if self.verbose: print "Dead time correction factors: ", CorrFactors
        #add third dimension to CorrFactors for broadcasting
        CorrFactors = np.reshape(CorrFactors,np.shape(CorrFactors)+(1,))
        return CorrFactors

    def loadAbsoluteSpectrum(self):
        '''
        extract the ARCONS measured spectrum of the spectrophotometric standard by breaking data into spectral cube
        and performing photometry (aper or psf) on each spectral frame
        '''
        if self.verbose:print "Making spectral cube"
        cubeDict = self.fluxFile.getSpectralCube(firstSec=self.startTime, integrationTime=self.intTime, weighted=True, fluxWeighted=False)
        cube= np.array(cubeDict['cube'], dtype=np.double)
        effIntTime= cubeDict['effIntTime']
        if self.verbose: print "median effective integration time in flux file cube = ", np.median(effIntTime)
        if self.verbose: print "cube shape ", np.shape(cube)
        if self.verbose: print "effIntTime shape ", np.shape(effIntTime)

        #add third dimension to effIntTime for broadcasting
        effIntTime = np.reshape(effIntTime,np.shape(effIntTime)+(1,))
        #put cube into counts/s in each pixel
        cube /= effIntTime
        #get dead time correction factors
        DTCorr = self.getDeadTimeCorrection(self.fluxFile)
        cube*=DTCorr #cube now in units of counts/s and corrected for dead time

        if self.plots and not 'figureHeader' in sys.modules:
            if self.verbose: print "Saving spectral frames as movie..."
            movieCube = np.zeros((self.nWvlBins,np.shape(cube)[0],np.shape(cube)[1]),dtype=float)
            for i in xrange(self.nWvlBins):
                movieCube[i,:,:] = cube[:,:,i]
            makeMovie(movieCube,frameTitles=self.binCenters,cbar=True,outName=self.plotSavePath+'FluxCal_Cube_%s.gif'%(self.objectName), normMin=0, normMax=50)
            if self.verbose: print "Movie saved in %s"%self.plotSavePath

        LCplot=False #light curve pop-ups not compatible with FLuxCal plotting 2/18/15
        #if self.photometry=='PSF': LCplot = False
        LC = LightCurve.LightCurve(verbose=self.verbose, showPlot=LCplot)

        self.fluxSpectrum=np.empty((self.nWvlBins),dtype=float)
        self.skySpectrum=np.zeros((self.nWvlBins),dtype=float)

        for i in xrange(self.nWvlBins):
            frame = cube[:,:,i]
            if self.verbose: print "%s photometry on frame %i of cube, central wvl = %f Angstroms"%(self.photometry,i,self.binCenters[i])
            if self.photometry == 'aperture':
                fDict = LC.performPhotometry(self.photometry,frame,[[self.centroidCol,self.centroidRow]],expTime=None,aper_radius = self.aperture, annulus_inner = self.annulusInner, annulus_outer = self.annulusOuter, interpolation="linear")
                self.fluxSpectrum[i] = fDict['flux']
                self.skySpectrum[i] = fDict['skyFlux']
                print "Sky estimate = ", fDict['skyFlux']
            else:
                fDict = LC.performPhotometry(self.photometry,frame,[[self.centroidCol,self.centroidRow]],expTime=None,aper_radius = self.aperture)
                self.fluxSpectrum[i] = fDict['flux']

        self.fluxSpectrum=self.fluxSpectrum/self.binWidths/self.collectingArea #spectrum now in counts/s/Angs/cm^2
        self.skySpectrum=self.skySpectrum/self.binWidths/self.collectingArea

        return self.fluxSpectrum, self.skySpectrum

    def loadRelativeSpectrum(self):
        self.fluxSpectra = [[[] for i in xrange(self.nCol)] for j in xrange(self.nRow)]
        self.fluxEffTime = [[[] for i in xrange(self.nCol)] for j in xrange(self.nRow)]
        for iRow in xrange(self.nRow):
            for iCol in xrange(self.nCol):
                count = self.fluxFile.getPixelCount(iRow,iCol)
                fluxDict = self.fluxFile.getPixelSpectrum(iRow,iCol,weighted=True,firstSec=0,integrationTime=-1)
                self.fluxSpectra[iRow][iCol],self.fluxEffTime[iRow][iCol] = fluxDict['spectrum'],fluxDict['effIntTime']
        self.fluxSpectra = np.array(self.fluxSpectra)
        self.fluxEffTime = np.array(self.fluxEffTime)
        DTCorr = self.getDeadTimeCorrection(self.fluxFile)
        #print "Bin widths = ",self.binWidths
        self.fluxSpectra = self.fluxSpectra/self.binWidths/self.fluxEffTime*DTCorr
        self.fluxSpectrum = self.calculateMedian(self.fluxSpectra) #find median of subtracted spectra across whole array
        return self.fluxSpectrum

    def loadSkySpectrum(self):
        self.skySpectra = [[[] for i in xrange(self.nCol)] for j in xrange(self.nRow)]
        self.skyEffTime = [[[] for i in xrange(self.nCol)] for j in xrange(self.nRow)]
        for iRow in xrange(self.nRow):
            for iCol in xrange(self.nCol):
                count = self.skyFile.getPixelCount(iRow,iCol)
                skyDict = self.skyFile.getPixelSpectrum(iRow,iCol,weighted=True,firstSec=0,integrationTime=-1)
                self.skySpectra[iRow][iCol],self.skyEffTime[iRow][iCol] = skyDict['spectrum'],skyDict['effIntTime']
        self.skySpectra = np.array(self.skySpectra)
        self.skyEffTime = np.array(self.skyEffTime)
        DTCorr = self.getDeadTimeCorrection(self.skyFile)
        self.skySpectra = self.skySpectra/self.binWidths/self.skyEffTime*DTCorr
        self.skySpectrum = self.calculateMedian(self.skySpectra) #find median of subtracted spectra across whole array
        return self.skySpectrum

    def loadStdSpectrum(self, objectName="G158-100"):
        #import the known spectrum of the calibrator and rebin to the histogram parameters given
        #must be imported into array with dtype float so division later does not have error
        std = MKIDStd.MKIDStd()
        a = std.load(objectName)
        a = std.countsToErgs(a) #convert std spectrum to ergs/s/Angs/cm^2 for BB fitting and cleaning
        self.stdWvls = np.array(a[:,0])
        self.stdFlux = np.array(a[:,1]) #std object spectrum in ergs/s/Angs/cm^2

        if self.plots:
            #create figure for plotting standard spectrum modifications
            self.stdFig = plt.figure()
            self.stdAx = self.stdFig.add_subplot(111)
            plt.xlim(3500,12000)
            plt.plot(self.stdWvls,self.stdFlux*1E15,linewidth=1,color='grey',alpha=0.75)

        convX_rev,convY_rev = self.cleanSpectrum(self.stdWvls,self.stdFlux)
        convX = convX_rev[::-1] #convolved spectrum comes back sorted backwards, from long wvls to low which screws up rebinning
        convY = convY_rev[::-1]
        #rebin cleaned spectrum to flat cal's wvlBinEdges
        newa = rebin(convX,convY,self.wvlBinEdges)
        rebinnedWvl = np.array(newa[:,0])
        rebinnedFlux = np.array(newa[:,1])
        if self.plots:
            #plot final resampled spectrum
            plt.plot(convX,convY*1E15,color='blue')
            plt.step(rebinnedWvl,rebinnedFlux*1E15,color = 'black',where='mid')
            plt.legend(['%s Spectrum'%self.objectName,'Blackbody Fit','Gaussian Convolved Spectrum','Rebinned Spectrum'],'upper right', numpoints=1)
            plt.xlabel(ur"Wavelength (\r{A})")
            plt.ylabel(ur"Flux (10$^{-15}$ ergs s$^{-1}$ cm$^{-2}$ \r{A}$^{-1}$)")
            plt.ylim(0.9*min(rebinnedFlux)*1E15, 1.1*max(rebinnedFlux)*1E15)
            plt.savefig(self.plotSavePath+'FluxCal_StdSpectrum_%s.eps'%self.objectName,format='eps')

        #convert standard spectrum back into counts/s/angstrom/cm^2
        newa = std.ergsToCounts(newa)
        self.binnedSpectrum = np.array(newa[:,1])

    def cleanSpectrum(self,x,y):
        ##=============== BB Fit to extend spectrum beyond 11000 Angstroms ==================
        fraction = 1.0/3.0
        nirX = np.arange(int(x[(1.0-fraction)*len(x)]),20000)
        T, nirY = fitBlackbody(x,y,fraction=fraction,newWvls=nirX,tempGuess=5600)

        if self.plots: plt.plot(nirX,nirY*1E15,linestyle='--',linewidth=2, color="black",alpha=0.5)

        extendedWvl = np.concatenate((x,nirX[nirX>max(x)]))
        extendedFlux = np.concatenate((y,nirY[nirX>max(x)]))
        ##======= Gaussian convolution to smooth std spectrum to MKIDs median resolution ========
        newX, newY = gaussianConvolution(extendedWvl,extendedFlux,xEnMin=0.005,xEnMax=6.0,xdE=0.001,fluxUnits = "lambda",r=self.r,plots=False)
        return newX, newY

    def calculateFactors(self):
        """
        Calculate the sensitivity spectrum: the weighting factors that correct the flat calibrated spectra to the real spectra

        For relative calibration:
        First subtract sky spectrum from ARCONS observed spectrum. Then take median of this spectrum as it should be identical
        across the array, assuming the flat cal has done its job. Then divide this into the known spectrum of the object.

        For absolute calibration:
        self.fluxSpectra already has sky subtraction included. Simply divide this spectrum into the known standard spectrum.
        """
        self.subtractedSpectrum = self.fluxSpectrum - self.skySpectrum
        self.subtractedSpectrum = np.array(self.subtractedSpectrum,dtype=float) #cast as floats so division does not fail later

        if self.method=='relative':
            normWvl = 5500 #Angstroms. Choose an arbitrary wvl to normalize the relative correction at
            ind = np.where(self.wvlBinEdges >= normWvl)[0][0]-1
            self.subtractedSpectrum = self.subtractedSpectrum/(self.subtractedSpectrum[ind]) #normalize
            self.binnedSpectrum = self.binnedSpectrum/(self.binnedSpectrum[ind]) #normalize treated Std spectrum while we are at it

        #Calculate FluxCal factors
        self.fluxFactors = self.binnedSpectrum/self.subtractedSpectrum

        #self.fluxFlags = np.zeros(np.shape(self.fluxFactors),dtype='int')
        self.fluxFlags = np.empty(np.shape(self.fluxFactors),dtype='int')
        self.fluxFlags.fill(pipelineFlags.fluxCal['good'])   #Initialise flag array filled with 'good' flags. JvE 5/1/2013.
        #set factors that will cause trouble to 1
        #self.fluxFlags[self.fluxFactors == np.inf] = 1
        self.fluxFlags[self.fluxFactors == np.inf] = pipelineFlags.fluxCal['infWeight']   #Modified to use flag dictionary - JvE 5/1/2013
        self.fluxFactors[self.fluxFactors == np.inf]=1.0
        self.fluxFlags[np.isnan(self.fluxFactors)] = pipelineFlags.fluxCal['nanWeight']   #Modified to use flag dictionary - JvE 5/1/2013
        self.fluxFactors[np.isnan(self.fluxFactors)]=1.0
        self.fluxFlags[self.fluxFactors <= 0]=pipelineFlags.fluxCal['LEzeroWeight']   #Modified to use flag dictionary - JvE 5/1/2013
        self.fluxFactors[self.fluxFactors <= 0]=1.0

    def calculateMedian(self, spectra):
        spectra2d = np.reshape(spectra,[self.nRow*self.nCol,self.nWvlBins])
        wvlMedian = np.empty(self.nWvlBins,dtype=float)
        for iWvl in xrange(self.nWvlBins):
            spectrum = spectra2d[:,iWvl]
            goodSpectrum = spectrum[spectrum != 0]#dead pixels need to be taken out before calculating medians
            wvlMedian[iWvl] = np.median(goodSpectrum)
        return wvlMedian

    def makePlots(self):
        """
        Output all debugging plots of ARCONS sky and object spectra, known calibrator spectrum, and sensitivity curve
        """
        scratchDir = os.getenv('MKID_PROC_PATH')
        fluxDir = self.plotSavePath
        fluxCalBase = 'FluxCal_%s'%self.objectName
        plotFileName = fluxCalBase+".pdf"
        fullFluxPlotFileName = os.path.join(fluxDir,plotFileName)

        #uncomment to make some plots for the paper. Proper formatting Will also require figureheader to be imported and for movie making to be turned off
        self.paperFig = plt.figure()
        self.paperAx = self.paperFig.add_subplot(111)
        plt.xlim(4000,11000)
        plt.plot(self.binCenters,self.fluxFactors,linewidth=3,color='black')
        plt.xlabel(ur"Wavelength (\r{A})")
        plt.ylabel(ur"Spectral Calibration Curve")
        plt.ylim(0,150)
        plt.savefig(self.plotSavePath+'FluxCal_Sensitivity_%s.eps'%self.objectName,format='eps')

        #save throughput as a .npz file that other code uses when making paper plots
        np.savez(self.plotSavePath+'%s_%s_throughput.npz'%(self.objectName.strip(),self.fluxTstamp),throughput=1.0/self.fluxFactors,wvls=self.binCenters)

        pp = PdfPages(fullFluxPlotFileName)
        #plt.rcParams['font.size'] = 2

        wvls = self.binCenters

        plt.figure()
        ax1 = plt.subplot(111)
        ax1.set_title('ARCONS median flat cal\'d flux in counts')
        plt.plot(wvls,self.fluxSpectrum)
        pp.savefig()

        plt.figure()
        ax2 = plt.subplot(111)
        ax2.set_title('ARCONS median flat cal\'d sky in counts')
        plt.plot(wvls,self.skySpectrum)
        pp.savefig()

        plt.figure()
        ax3 = plt.subplot(111)
        ax3.set_title('Flux data minus sky in counts')
        plt.plot(wvls,self.subtractedSpectrum)
        pp.savefig()

        plt.figure()
        ax4 = plt.subplot(111)
        ax4.set_title('Std Spectrum of %s'%(self.objectName))
        plt.plot(self.stdWvls,self.stdFlux)
        pp.savefig()

        plt.figure()
        ax5 = plt.subplot(111)
        ax5.set_title('Binned Std Spectrum')
        plt.plot(wvls,self.binnedSpectrum)
        pp.savefig()

        plt.figure()
        ax6 = plt.subplot(111)
        ax6.set_title('Median Sensitivity Spectrum')
        ax6.set_xlim((3500,12000))
        #ax6.set_ylim((0,5))
        plt.plot(wvls,self.fluxFactors)
        pp.savefig()

        plt.figure()
        ax7 = plt.subplot(111)
        ax7.set_title('1/Sensitivity (Throughput)')
        ax7.set_xlim((3500,12000))
        ax7.set_ylim((0,.04))
        plt.plot(wvls,1.0/self.fluxFactors)
        pp.savefig()

        plt.figure()
        ax8 = plt.subplot(111)
        ax8.set_title('Flux Cal\'d ARCONS Spectrum of Std')
        plt.plot(wvls,self.fluxFactors*self.subtractedSpectrum)
        pp.savefig()
        pp.close()
        print "Saved Flux Cal plots to %s"%(fullFluxPlotFileName)

    def writeFactors(self,fluxCalFileName):
        """
        Write flux cal weights to h5 file
        """

        if os.path.isabs(fluxCalFileName) == True:
            fullFluxCalFileName = fluxCalFileName
        else:
            scratchDir = os.getenv('MKID_PROC_PATH')
            fluxDir = os.path.join(scratchDir,'fluxCalSolnFiles')
            fullFluxCalFileName = os.path.join(fluxDir,fluxCalFileName)

        try:
            fluxCalFile = tables.openFile(fullFluxCalFileName,mode='w')
        except:
            print 'Error: Couldn\'t create flux cal file, ',fullFluxCalFileName
            return

        calgroup = fluxCalFile.createGroup(fluxCalFile.root,'fluxcal','Table of flux calibration weights by wavelength')
        caltable = tables.Array(calgroup,'weights',object=self.fluxFactors,title='Flux calibration Weights indexed by wavelengthBin')
        flagtable = tables.Array(calgroup,'flags',object=self.fluxFlags,title='Flux cal flags indexed by wavelengthBin. 0 is Good')
        bintable = tables.Array(calgroup,'wavelengthBins',object=self.wvlBinEdges,title='Wavelength bin edges corresponding to third dimension of weights array')
        fluxCalFile.flush()
        fluxCalFile.close()
        print "Finished Flux Cal, written to %s"%(fullFluxCalFileName)

    def cleanSpectrum_old(self,x,y,objectName):
        '''
        function to take high resolution spectrum of standard star, extend IR coverage with
        an exponential tail, then rebin down to ARCONS resolution. This function has since been
        deprecated with the current cleanSpectrum which uses a BB fit to extend IR coverage,
        and does the rebinning using a gaussian convolution. This is left in for reference.
        '''

        #locations and widths of absorption features in Angstroms
        #features = [3890,3970,4099,4340,4860,6564,6883,7619]
        #widths = [50,50,50,50,50,50,50,50]
        #for i in xrange(len(features)):
        #    #check for absorption feature in std spectrum
        #    ind = np.where((x<(features[i]+15)) & (x>(features[i]-15)))[0]
        #    if len(ind)!=0:
        #        ind = ind[len(ind)/2]
        #    #if feature is found (flux is higher on both sides of the specified wavelength where the feature should be)
        #    if y[ind]<y[ind+1] and y[ind]<y[ind-1]:
        #        #cut out width[i] around feature[i]
        #        inds = np.where((x >= features[i]+widths[i]) | (x <= features[i]-widths[i]))
        #        x = x[inds]
        #        y = y[inds]

        #fit a tail to the end of the spectrum to interpolate out to desired wavelength in angstroms
        fraction = 3.0/4.0
        newx = np.arange(int(x[fraction*len(x)]),20000)

        slopeguess = (np.log(y[-1])-np.log(y[fraction*len(x)]))/(x[-1]-x[fraction*len(x)])
        print "Guess at exponential slope is %f"%(slopeguess)
        guess_a, guess_b, guess_c = float(y[fraction*len(x)]), x[fraction*len(x)], slopeguess
        guess = [guess_a, guess_b, guess_c]

        fitx = x[fraction*len(x):]
        fity = y[fraction*len(x):]

        exp_decay = lambda fx, A, x0, t: A * np.exp((fx-x0) * t)

        params, cov = curve_fit(exp_decay, fitx, fity, p0=guess, maxfev=2000)
        A, x0, t= params
        print "A = %s\nx0 = %s\nt = %s\n"%(A, x0, t)
        best_fit = lambda fx: A * np.exp((fx-x0)*t)

        calcx = np.array(newx,dtype=float)
        newy = best_fit(calcx)

        #func = interpolate.splrep(x[fration*len(x):],y[fraction*len(x):],s=smooth)
        #newx = np.arange(int(x[fraction*len(x)]),self.wvlBinEdges[-1])
        #newy = interpolate.splev(newx,func)

        wl = np.concatenate((x,newx[newx>max(x)]))
        flux = np.concatenate((y,newy[newx>max(x)]))

        #new method, rebin data to grid of wavelengths generated from a grid of evenly spaced energy bins
        #R=7.0 at 4500
        #R=E/dE -> dE = R/E
        dE = 0.3936 #eV
        start = 1000 #Angs
        stop = 20000 #Angs
        enBins = ObsFile.makeWvlBins(dE,start,stop)
        rebinned = rebin(wl,flux,enBins)
        re_wl = rebinned[:,0]
        re_flux = rebinned[:,1]
        #plt.plot(re_wl,re_flux,color='r')

        re_wl = re_wl[np.isnan(re_flux)==False]
        re_flux = re_flux[np.isnan(re_flux)==False]

        start1 = self.wvlBinEdges[0]
        stop1 = self.wvlBinEdges[-1]
        #regrid downsampled data

        new_wl = np.arange(start1,stop1)

        #print re_wl
        #print re_flux
        #print new_wl

        #weight=1.0/(re_flux)**(2/1.00)
        print len(re_flux)
        weight = np.ones(len(re_flux))
        #decrease weights near peak
        ind = np.where(re_flux == max(re_flux))[0]
        weight[ind] = 0.3
        for p in [1,2,3]:
            if p==1:
                wt = 0.3
            elif p==2:
                wt = 0.6
            elif p==3:
                wt = 0.7
            try:
                weight[ind+p] = wt
            except IndexError:
                 pass
            try:
                 if ind-p >= 0:
                     weight[ind-p] = wt
            except IndexError:
                pass
        weight[-4:] = 1.0
        #weight = [0.7,1,0.3,0.3,0.5,0.7,1,1,1]
        #print len(weight)
        #weight = re_flux/min(re_flux)
        #weight = 1.0/weight
        #weight = weight/max(weight)
        #print weight
        f = interpolate.splrep(re_wl,re_flux,w=weight,k=3,s=max(re_flux)**1.71)
        new_flux = interpolate.splev(new_wl,f,der=0)
        return new_wl, new_flux

if __name__ == '__main__':
    try:
        paramFile = sys.argv[1]
    except:
        paramFile = '/home/srmeeker/ARCONS-pipeline/params/fluxCal.dict'
    fc = FluxCal(paramFile, plots=True, verbose=True)


"""
===============================
Nearest Centroid Classification
===============================

Sample usage of Nearest Centroid classification.
It will plot the decision boundaries for each class.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn import datasets
from sklearn.neighbors import NearestCentroid

n_neighbors = 15

# import some data to play with
iris = datasets.load_iris()
X = iris.data[:, :2]  # we only take the first two features. We could
                      # avoid this ugly slicing by using a two-dim dataset
y = iris.target

h = .02  # step size in the mesh

# Create color maps
cmap_light = ListedColormap(['#FFAAAA', '#AAFFAA', '#AAAAFF'])
cmap_bold = ListedColormap(['#FF0000', '#00FF00', '#0000FF'])

for shrinkage in [None, 0.1]:
    # we create an instance of Neighbours Classifier and fit the data.
    clf = NearestCentroid(shrink_threshold=shrinkage)
    clf.fit(X, y)
    y_pred = clf.predict(X)
    print(shrinkage, np.mean(y == y_pred))
    # Plot the decision boundary. For that, we will assign a color to each
    # point in the mesh [x_min, m_max]x[y_min, y_max].
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                         np.arange(y_min, y_max, h))
    Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    plt.figure()
    plt.pcolormesh(xx, yy, Z, cmap=cmap_light)

    # Plot also the training points
    plt.scatter(X[:, 0], X[:, 1], c=y, cmap=cmap_bold)
    plt.title("3-Class classification (shrink_threshold=%r)"
              % shrinkage)
    plt.axis('tight')

plt.show()


import glob
import os

import pandas as pd
import openmc
import pytest

from tests.testing_harness import PyAPITestHarness

class DiffTallyTestHarness(PyAPITestHarness):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Set settings explicitly
        self._model.settings.batches = 3
        self._model.settings.inactive = 0
        self._model.settings.particles = 100
        self._model.settings.source = openmc.Source(space=openmc.stats.Box(
            [-160, -160, -183], [160, 160, 183]))
        self._model.settings.temperature['multipole'] = True

        filt_mats = openmc.MaterialFilter((1, 3))
        filt_eout = openmc.EnergyoutFilter((0.0, 0.625, 20.0e6))

        # We want density derivatives for both water and fuel to get coverage
        # for both fissile and non-fissile materials.
        d1 = openmc.TallyDerivative(derivative_id=1)
        d1.variable = 'density'
        d1.material = 3
        d2 = openmc.TallyDerivative(derivative_id=2)
        d2.variable = 'density'
        d2.material = 1

        # O-16 is a good nuclide to test against because it is present in both
        # water and fuel.  Some routines need to recognize that they have the
        # perturbed nuclide but not the perturbed material.
        d3 = openmc.TallyDerivative(derivative_id=3)
        d3.variable = 'nuclide_density'
        d3.material = 1
        d3.nuclide = 'O16'

        # A fissile nuclide, just for good measure.
        d4 = openmc.TallyDerivative(derivative_id=4)
        d4.variable = 'nuclide_density'
        d4.material = 1
        d4.nuclide = 'U235'

        # Temperature derivatives.
        d5 = openmc.TallyDerivative(derivative_id=5)
        d5.variable = 'temperature'
        d5.material = 1

        derivs = [d1, d2, d3, d4, d5]

        # Cover the flux score.
        for i in range(5):
            t = openmc.Tally()
            t.scores = ['flux']
            t.filters = [filt_mats]
            t.derivative = derivs[i]
            self._model.tallies.append(t)

        # Cover supported scores with a collision estimator.
        for i in range(5):
            t = openmc.Tally()
            t.scores = ['total', 'absorption', 'scatter', 'fission', 'nu-fission']
            t.filters = [filt_mats]
            t.nuclides = ['total', 'U235']
            t.derivative = derivs[i]
            self._model.tallies.append(t)

        # Cover an analog estimator.
        for i in range(5):
            t = openmc.Tally()
            t.scores = ['absorption']
            t.filters = [filt_mats]
            t.estimator = 'analog'
            t.derivative = derivs[i]
            self._model.tallies.append(t)

        # Energyout filter and total nuclide for the density derivatives.
        for i in range(2):
            t = openmc.Tally()
            t.scores = ['nu-fission', 'scatter']
            t.filters = [filt_mats, filt_eout]
            t.nuclides = ['total', 'U235']
            t.derivative = derivs[i]
            self._model.tallies.append(t)

        # Energyout filter without total nuclide for other derivatives.
        for i in range(2, 5):
            t = openmc.Tally()
            t.scores = ['nu-fission', 'scatter']
            t.filters = [filt_mats, filt_eout]
            t.nuclides = ['U235']
            t.derivative = derivs[i]
            self._model.tallies.append(t)

    def _get_results(self):
        # Read the statepoint and summary files.
        statepoint = glob.glob(os.path.join(os.getcwd(), self._sp_name))[0]
        sp = openmc.StatePoint(statepoint)

        # Extract the tally data as a Pandas DataFrame.
        df = pd.DataFrame()
        for t in sp.tallies.values():
            df = df.append(t.get_pandas_dataframe(), ignore_index=True)

        # Extract the relevant data as a CSV string.
        cols = ('d_material', 'd_nuclide', 'd_variable', 'score', 'mean',
                'std. dev.')
        return df.to_csv(None, columns=cols, index=False, float_format='%.7e')

def test_diff_tally():
    harness = DiffTallyTestHarness('statepoint.3.h5')
    harness.main()


#!/usr/bin/env python

# Generate overlay images in PNG format with transparancy which can be
# used to label Splotch frames.  This script can be called as a
# standalone program, see below for details.  To label an entire
# directory of Splotch frames, use the driver script <splotchLabelFrames.sh>.
#
#                                         (Klaus Reuter, RZG, Sep 2011)

def splotchColormap(time=-1.0,               # for time>0, a time stamp is printed in the upper left corner
                    redshift=-1.0,           # for redshift>0, a redshift stamp is printed
                    valMin=0.1,              # minimum value for the log colorscale
                    valMax=1.e4,             # maximum value for the log colorscale
                    outfile="overlay.png",   # default file name of the overlay to be created
                    xinches=12,              # width of the image  | at 100 DPI, this corresponds to
                    yinches=8,               # height of the image | the dimensions 1200x800
                    myFontSize="large",
                    myFontColor="white",
                    putMinerva=False):       # place the MPG minerva logo in the top right corner

   # import necessary modules
   import numpy as np
   from matplotlib import pyplot
   import matplotlib as mpl
   from subprocess import call
   from math import pow

   # *** set font properties for annotations ***
   fprops=mpl.font_manager.FontProperties()
   fprops.set_size(myFontSize)
   #fprops.set_weight("bold")

   # *** set up the matplotlib colormap based on a Splotch colormap ***

   #$ cat OldSplotch.pal
   #OldSplotch
   #0100
   #3
   #  0   0 255
   #128 255 128
   #255   0   0

   # See <http://matplotlib.sourceforge.net/api/colors_api.html>
   #  to understand what's going on ...
   # <OldSplotch.pal> corresponds to:
   OldSplotch = {'red':   ((0.0, 0.0, 0.0), (0.5, 0.5, 0.5), (1.0, 1.0, 1.0)),
                 'green': ((0.0, 0.0, 0.0), (0.5, 1.0, 1.0), (1.0, 0.0, 0.0)),
                 'blue':  ((0.0, 1.0, 1.0), (0.5, 0.5, 0.5), (1.0, 0.0, 0.0))}
   colormap = mpl.colors.LinearSegmentedColormap('colormap', OldSplotch)
   # TODO implement a reader for Splotch palette files

   # *** set up the figure ***
   fig = pyplot.figure(figsize=(xinches,yinches))

   # *** set up the colorbar ***
   ax1  = fig.add_axes([0.90, 0.05, 0.02, 0.5])
   norm = mpl.colors.LogNorm(vmin=valMin, vmax=valMax)
   form = mpl.ticker.LogFormatterMathtext()
   cb1  = mpl.colorbar.ColorbarBase(ax1, cmap=colormap, norm=norm,
                                    format=form, orientation='vertical')
   # manipulate the style of the ticklabels, which requires a loop
   for tl in cb1.ax.get_yticklabels():
      tl.set_fontsize(myFontSize)
      tl.set_color(myFontColor)
   cb1.set_label('Temperature [K]', fontproperties=fprops, color=myFontColor)

   # *** set up the time/redshift variable ***
   if (time>=0.0):
      timeString="age of universe=%.3f" % (time, )
      timeString=timeString+" Gyr"
      pyplot.figtext(x=0.025, y=0.950, s=timeString, fontdict=None,
                     fontproperties=fprops, color=myFontColor)
   #
   if (redshift>0):
      timeString="redshift=%.3f" % (redshift, )
      pyplot.figtext(x=0.025, y=0.910, s=timeString, fontdict=None,
                     fontproperties=fprops, color=myFontColor)

   # Minerva needs an intermediate call of the ImageMagick tools
   if putMinerva:
      plotFile="./splotchColormapTmp.png"
   else:
      plotFile=outfile

   # *** finally, plot the image and write it to a png file ***
   pyplot.plot()
   F=pyplot.gcf()
   myDPI=100
   F.savefig(plotFile, transparent=True, dpi=myDPI)

   # *** put a logo (e.g. MPG Minerva) on top using ImageMagick convert ***
   if putMinerva:
      minervaFile="__INSERT_VALID_PATH__/minerva-white-96.png"
      xoffset=str(int( (xinches*myDPI)*0.895 ))
      yoffset=str(int( (yinches*myDPI)*0.005 ))
      #print (xoffset, yoffset)
      convertCommand="/usr/bin/env convert "+plotFile+" "+minervaFile+" -geometry +"+xoffset+"+"+yoffset+" -composite -format png "+outfile
      call(convertCommand, shell=True)

   # *** END SplotchColormap() ***

#
# *** Allow this Python module to be run as a standalone script. ***
#
if __name__ == "__main__":
   import sys
   import getopt
   #
   try:
      opts, args = getopt.getopt(sys.argv[1:],
                        "t:r:c:d:o:",  # the "-" options, below are the "--" options
                        ["time=", "redshift=", "colormin=", "colormax=", "outfile="])
   except getopt.GetoptError, err:
      print str(err)
      sys.exit(2)
   #
   myOutFile  = "overlay.png"
   myTime     = -1.0
   myRedshift = -1.0
   myMinVal   = 1
   myMaxVal   = 100
   #
   for o, a in opts:
      # print (o,a)
      if   o in ("-t", "--time"):
         myTime     = float(a)
      elif o in ("-r", "--redshift"):
         myRedshift = float(a)
      elif o in ("-c", "--colormin"):
         myMinVal   = pow(10.0, float(a))
      elif o in ("-d", "--colormax"):
         myMaxVal   = pow(10.0, float(a))
      elif o in ("-o", "--outfile"):
         myOutFile  = a
      else:
          assert False, "unhandled option"
   #
   splotchColormap(outfile=myOutFile,
                   time=myTime,
                   redshift=myRedshift,
                   valMin=myMinVal,
                   valMax=myMaxVal)

# EOF


import numpy as np
import time
import matplotlib.pyplot as plt
import random
# get gekko package with:
#   pip install gekko
from gekko import GEKKO
import pandas as pd

# import data
data = pd.read_csv('data.txt')
tm = data['Time (sec)'].values
Q1s = data[' Heater 1'].values
Q2s = data[' Heater 2'].values
T1s = data[' Temperature 1'].values
T2s = data[' Temperature 2'].values

#########################################################
# Initialize Model as Estimator
#########################################################
m = GEKKO(name='tclab-mhe')
#m.server = 'http://127.0.0.1' # if local server is installed

# 120 second time horizon, 40 steps
m.time = tm

# Parameters to Estimate
K1 = m.FV(value=0.5)
K1.STATUS = 1
K1.FSTATUS = 0
K1.LOWER = 0.1
K1.UPPER = 1.0

K2 = m.FV(value=0.3)
K2.STATUS = 1
K2.FSTATUS = 0
K2.LOWER = 0.1
K2.UPPER = 1.0

K3 = m.FV(value=0.1)
K3.STATUS = 1
K3.FSTATUS = 0
K3.LOWER = 0.0001
K3.UPPER = 1.0

tau12 = m.FV(value=150)
tau12.STATUS = 1
tau12.FSTATUS = 0
tau12.LOWER = 50.0
tau12.UPPER = 250

tau3 = m.FV(value=15)
tau3.STATUS = 0
tau3.FSTATUS = 0
tau3.LOWER = 10
tau3.UPPER = 20

# Measured inputs
Q1 = m.MV(value=0)
Q1.FSTATUS = 1 # measured
Q1.value = Q1s

Q2 = m.MV(value=0)
Q2.FSTATUS = 1 # measured
Q2.value = Q2s

# Ambient temperature
Ta = m.Param(value=23.0) # degC

# State variables
TH1 = m.SV(value=T1s[0])
TH2 = m.SV(value=T2s[0])

# Measurements for model alignment
TC1 = m.CV(value=T1s)
TC1.STATUS = 1     # minimize error between simulation and measurement
TC1.FSTATUS = 1    # receive measurement
TC1.MEAS_GAP = 0.1 # measurement deadband gap

TC2 = m.CV(value=T1s[0])
TC2.STATUS = 1     # minimize error between simulation and measurement
TC2.FSTATUS = 1    # receive measurement
TC2.MEAS_GAP = 0.1 # measurement deadband gap
TC2.value = T2s

# Heat transfer between two heaters
DT = m.Intermediate(TH2-TH1)

# Empirical correlations
m.Equation(tau12 * TH1.dt() + (TH1-Ta) == K1*Q1 + K3*DT)
m.Equation(tau12 * TH2.dt() + (TH2-Ta) == K2*Q2 - K3*DT)
m.Equation(tau3 * TC1.dt()  + TC1 == TH1)
m.Equation(tau3 * TC2.dt()  + TC2 == TH2)

# Global Options
m.options.IMODE   = 5 # MHE
m.options.EV_TYPE = 2 # Objective type
m.options.NODES   = 3 # Collocation nodes
m.options.SOLVER  = 3 # IPOPT
m.options.COLDSTART = 0 # COLDSTART on first cycle

# Predict Parameters and Temperatures
# use remote=False for local solve
m.solve()

# Create plot
plt.figure(figsize=(10,7))

ax=plt.subplot(2,1,1)
ax.grid()
plt.plot(tm,T1s,'ro',label=r'$T_1$ measured')
plt.plot(tm,TC1.value,'k-',label=r'$T_1$ predicted')
plt.plot(tm,T2s,'bx',label=r'$T_2$ measured')
plt.plot(tm,TC2.value,'k--',label=r'$T_2$ predicted')
plt.ylabel('Temperature (degC)')
plt.legend(loc=2)
ax=plt.subplot(2,1,2)
ax.grid()
plt.plot(tm,Q1s,'r-',label=r'$Q_1$')
plt.plot(tm,Q2s,'b:',label=r'$Q_2$')
plt.ylabel('Heaters')
plt.xlabel('Time (sec)')
plt.legend(loc='best')

# Print optimal values
print('K1: ' + str(K1.newval))
print('K2: ' + str(K2.newval))
print('K3: ' + str(K3.newval))
print('tau12: ' + str(tau12.newval))
print('tau3: ' + str(tau3.newval))

# Save figure
plt.savefig('tclab_estimation.png')
plt.show()


# yellowbrick.contrib.scatter
# Implements a 2d scatter plot for feature analysis.
#
# Author:   Nathan Danielsen
# Created:  Fri Feb 26 19:40:00 2017 -0400
#
# Copyright (C) 2017 The scikit-yb developers
# For license information, see LICENSE.txt
#
# ID: scatter.py [a89633e] benjamin@bengfort.com $
"""
Implements a 2D scatter plot for feature analysis.
"""

##########################################################################
# Imports
##########################################################################

import itertools
import numpy as np

from yellowbrick.features.base import DataVisualizer
from yellowbrick.utils import is_dataframe, is_structured_array
from yellowbrick.utils import has_ndarray_int_columns
from yellowbrick.exceptions import YellowbrickValueError
from yellowbrick.style.colors import resolve_colors

##########################################################################
# Quick Methods
##########################################################################

def scatterviz(
    X,
    y=None,
    ax=None,
    features=None,
    classes=None,
    color=None,
    colormap=None,
    markers=None,
    alpha=1.0,
    **kwargs
):
    """Displays a bivariate scatter plot.

    This helper function is a quick wrapper to utilize the ScatterVisualizer
    (Transformer) for one-off analysis.

    Parameters
    ----------

    X : ndarray or DataFrame of shape n x m
        A matrix of n instances with m features

    y : ndarray or Series of length n, default: None
        An array or series of target or class values

    ax : matplotlib axes, default: None
        The axes to plot the figure on.

    features : list of strings, default: None
        The names of two features or columns.
        More than that will raise an error.

    classes : list of strings, default: None
        The names of the classes in the target

    color : list or tuple of colors, default: None
        Specify the colors for each individual class

    colormap : string or matplotlib cmap, default: None
        Sequential colormap for continuous target

    markers : iterable of strings, default: ,+o*vhd
        Matplotlib style markers for points on the scatter plot points

    alpha : float, default: 1.0
        Specify a transparency where 1 is completely opaque and 0 is completely
        transparent. This property makes densely clustered points more visible.

    Returns
    -------
    viz : ScatterVisualizer
        Returns the fitted, finalized visualizer
    """
    # Instantiate the visualizer
    visualizer = ScatterVisualizer(
        ax=ax,
        features=features,
        classes=classes,
        color=color,
        colormap=colormap,
        markers=markers,
        alpha=alpha,
        **kwargs
    )

    # Fit and transform the visualizer (calls draw)
    visualizer.fit(X, y, **kwargs)
    visualizer.transform(X)

    # Return the visualizer object
    return visualizer

##########################################################################
# Static ScatterVisualizer Visualizer
##########################################################################

class ScatterVisualizer(DataVisualizer):
    """
    ScatterVisualizer is a bivariate feature data visualization algorithm that
    plots using the Cartesian coordinates of each point.

        Parameters
        ----------

        ax : a matplotlib plot, default: None
            The axis to plot the figure on.

        x : string, default: None
            The feature name that corresponds to a column name or index postion
            in the matrix that will be plotted against the x-axis

        y : string, default: None
            The feature name that corresponds to a column name or index postion
            in the matrix that will be plotted against the y-axis

        features : a list of two feature names to use, default: None
            List of two features that correspond to the columns in the array.
            The order of the two features correspond to X and Y axes on the
            graph. More than two feature names or columns will raise an error.
            If a DataFrame is passed to fit and features is None, feature names
            are selected that are the columns of the DataFrame.

        classes : a list of class names for the legend, default: None
            If classes is None and a y value is passed to fit then the classes
            are selected from the target vector.

        color : optional list or tuple of colors to colorize points, default: None
            Use either color to colorize the points on a per class basis or
            colormap to color them on a continuous scale.

        colormap : optional string or matplotlib cmap to colorize points, default: None
            Use either color to colorize the points on a per class basis or
            colormap to color them on a continuous scale.

        markers : iterable of strings, default: ,+o*vhd
            Matplotlib style markers for points on the scatter plot points

        alpha : float, default: 1.0
            Specify a transparency where 1 is completely opaque and 0 is completely
            transparent. This property makes densely clustered points more visible.

        kwargs : keyword arguments passed to the super class.

        These parameters can be influenced later on in the visualization
        process, but can and should be set as early as possible.
    """

    def __init__(
        self,
        ax=None,
        x=None,
        y=None,
        features=None,
        classes=None,
        color=None,
        colormap=None,
        markers=None,
        alpha=1.0,
        **kwargs
    ):
        """
        Initialize the base scatter with many of the options required in order
        to make the visualization work.
        """
        super(ScatterVisualizer, self).__init__(
            ax=ax,
            features=features,
            classes=classes,
            color=color,
            colormap=colormap,
            **kwargs
        )

        self.x = x
        self.y = y

        self.alpha = alpha

        self.markers = itertools.cycle(
            kwargs.pop("markers", (",", "+", "o", "*", "v", "h", "d"))
        )

        self.color = color
        self.colormap = colormap

        if self.x is not None and self.y is not None and self.features is not None:
            raise YellowbrickValueError("Please specify x,y or features, not both.")

        if self.x is not None and self.y is not None and self.features is None:
            self.features = [self.x, self.y]

        # Ensure with init that features doesn't have more than two features
        if features is not None:
            if len(features) != 2:
                raise YellowbrickValueError(
                    "ScatterVisualizer only accepts two features."
                )

    def fit(self, X, y=None, **kwargs):
        """
        The fit method is the primary drawing input for the parallel coords
        visualization since it has both the X and y data required for the
        viz and the transform method does not.

        Parameters
        ----------
        X : ndarray or DataFrame of shape n x m
            A matrix of n instances with 2 features

        y : ndarray or Series of length n
            An array or series of target or class values

        kwargs : dict
            Pass generic arguments to the drawing method

        Returns
        -------
        self : instance
            Returns the instance of the transformer/visualizer
        """
        _, ncols = X.shape

        # NOTE: Do not call super for this class, it conflicts with the fit.
        # Setting these variables is similar to the old behavior of DataVisualizer.
        # TODO: refactor to make use of the new DataVisualizer functionality
        self.features_ = self.features
        self.classes_ = self.classes

        if ncols == 2:
            X_two_cols = X
            if self.features_ is None:
                self.features_ = ["Feature One", "Feature Two"]

        # Handle the feature names if they're None.
        elif self.features_ is not None and is_dataframe(X):
            X_two_cols = X[self.features_].values

        # handle numpy named/ structured array
        elif self.features_ is not None and is_structured_array(X):
            X_selected = X[self.features_]
            X_two_cols = X_selected.copy().view(
                (np.float64, len(X_selected.dtype.names))
            )

        # handle features that are numeric columns in ndarray matrix
        elif self.features_ is not None and has_ndarray_int_columns(self.features_, X):
            f_one, f_two = self.features_
            X_two_cols = X[:, [int(f_one), int(f_two)]]

        else:
            raise YellowbrickValueError(
                """
                ScatterVisualizer only accepts two features, please
                explicitly set these two features in the init kwargs or
                pass a matrix/ dataframe in with only two columns."""
            )

        # Store the classes for the legend if they're None.
        if self.classes_ is None:
            # TODO: Is this the most efficient method?
            self.classes_ = [str(label) for label in np.unique(y)]

        # Draw the instances
        self.draw(X_two_cols, y, **kwargs)

        # Fit always returns self.
        return self

    def draw(self, X, y, **kwargs):
        """Called from the fit method, this method creates a scatter plot that
        draws each instance as a class or target colored point, whose location
        is determined by the feature data set.
        """
        # Set the axes limits
        self.ax.set_xlim([-1, 1])
        self.ax.set_ylim([-1, 1])

        # set the colors
        color_values = resolve_colors(
            n_colors=len(self.classes_), colormap=self.colormap, colors=self.color
        )

        colors = dict(zip(self.classes_, color_values))

        # Create a data structure to hold the scatter plot representations
        to_plot = {}
        for kls in self.classes_:
            to_plot[kls] = [[], []]

        # Add each row of the data set to to_plot for plotting
        # TODO: make this an independent function for override
        for i, row in enumerate(X):
            row_ = np.repeat(np.expand_dims(row, axis=1), 2, axis=1)
            x_, y_ = row_[0], row_[1]
            kls = self.classes_[y[i]]

            to_plot[kls][0].append(x_)
            to_plot[kls][1].append(y_)

        # Add the scatter plots from the to_plot function
        # TODO: store these plots to add more instances to later
        # TODO: make this a separate function
        for i, kls in enumerate(self.classes_):
            self.ax.scatter(
                to_plot[kls][0],
                to_plot[kls][1],
                marker=next(self.markers),
                color=colors[kls],
                label=str(kls),
                alpha=self.alpha,
                **kwargs
            )

        self.ax.axis("equal")

    def finalize(self, **kwargs):
        """
        Adds a title and a legend and ensures that the axis labels are set as
        the feature names being visualized.

        Parameters
        ----------
        kwargs: generic keyword arguments.

        Notes
        -----
        Generally this method is called from show and not directly by the user.
        """
        # Divide out the two features
        feature_one, feature_two = self.features_

        # Set the title
        self.set_title(
            "Scatter Plot: {0} vs {1}".format(str(feature_one), str(feature_two))
        )
        # Add the legend
        self.ax.legend(loc="best")
        self.ax.set_xlabel(str(feature_one))
        self.ax.set_ylabel(str(feature_two))

# Alias for ScatterViz
ScatterViz = ScatterVisualizer


import matplotlib
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(111)
patch1 = matplotlib.patches.Circle(
    [0.5,0.5],0.05
)
patch2 = matplotlib.patches.Rectangle(
    [0.3,0.3],0.4, 0.4, alpha=0.5,
    fill=False, edgecolor='black',
    linestyle = '--'
)
arrow1 = matplotlib.patches.Arrow(
    0, 0.5,0.45,0, width=0.05,
    color='black'
)
arrow2 = matplotlib.patches.Arrow(
    0.55, 0.5,0.45,0, width=0.05,
    color='black'
)
line1 = matplotlib.lines.Line2D(
    [0.5,0.5], [0,0.45],
    linestyle='--', color='black'
)
text1 = matplotlib.text.Text(
    0, 0.45, '$n_{A0}$\n$V_0$\n$U_A=0$'
)
text2 = matplotlib.text.Text(
    0.8, 0.45, '$n_{A1}$\n$V_1$\n$U_{A1}$'
)
for artist in [
    patch1,patch2,arrow1,arrow2,
    line1,text1,text2
]:
    ax.add_artist(artist)
ax.set_frame_on(False)
ax.set_axis_off()
ax.set_aspect(1.0)

fig.

fig = plt.figure()
ax = fig.add_subplot(111)
patch1 = matplotlib.patches.Circle(
    [0.5,0.5],0.05
)
patch2 = matplotlib.patches.Rectangle(
    [0.3,0.3],0.4, 0.4, alpha=0.5,
    fill=False, edgecolor='black',
    linestyle = '--'
)
arrow1 = matplotlib.patches.Arrow(
    0, 0.5,0.45,0, width=0.05,
    color='black'
)
arrow2 = matplotlib.patches.Arrow(
    0.55, 0.5,0.45,0, width=0.05,
    color='black'
)
arrow3 = matplotlib.patches.Arrow(
    0.5, 0.0, 0,0.45, width=0.05,
    color='black'
)
text1 = matplotlib.text.Text(
    0, 0.45, '$n_{A0}$\n$V_0$\n$U_A=0$'
)
text2 = matplotlib.text.Text(
    0.8, 0.45, '$n_{A1}$\n$V_1$\n$U_{A1}$'
)
text3 = matplotlib.text.Text(
    0.55, 0.1, '$n_{Ar}$\n$V_r$'
)
for artist in [
    patch1,patch2,arrow1,arrow2,
    arrow3,text1,text2,text3
]:
    ax.add_artist(artist)
ax.set_frame_on(False)
ax.set_axis_off()
ax.set_aspect(1.0)