# -*- coding: utf-8 -*-
"""
Created on Mon Apr 24 14:13:47 2017

@author: ishort
"""

import math
import Useful
import ToolBox
#import numpy
#JB#
#from matplotlib.pyplot import plot, title, show, scatter
#storage for fits (not all may be used)
uw = []
uwa = []
uwb = []
uwStage = []
uwbStage = []
uwu = []
uwl = []
uua=[]
uub=[]
"""
#a function to create a cubic function fit extrapolation
def cubicFit(x,y):
    coeffs = numpy.polyfit(x,y,3)
    #returns an array of coefficents for the cubic fit of the form
    #Ax^3 + Bx^2 + Cx + D as [A,B,C,D]
    return coeffs
#this will work for any number of data points!
def valueFromFit(fit,x):
    #return the value y for a given fit, at point x
    return (fit[0]*(x**3)+fit[1]*(x**2)+fit[2]*x+fit[3])
#holds the five temperature at which we have partition function data
"""
masterTemp = [130, 500, 3000, 8000, 10000]
#JB#

#def levelPops(lam0In, logNStage, chiL, log10UwStage, gwL, numDeps, temp):
def levelPops(lam0In, logNStage, chiL, logUw, gwL, numDeps, temp):
    """ Returns depth distribution of occupation numbers in lower level of b-b transition,

// Input parameters:
// lam0 - line centre wavelength in nm
// logNStage - log_e density of absorbers in relevent ion stage (cm^-3)
// logFlu - log_10 oscillator strength (unitless)
// chiL - energy of lower atomic E-level of b-b transition in eV
// Also needs atsmopheric structure information:
// numDeps
// temp structure """

    c = Useful.c()
    logC = Useful.logC()
    k = Useful.k()
    logK = Useful.logK()
    logH = Useful.logH()
    logEe = Useful.logEe()
    logMe = Useful.logMe()

    ln10 = math.log(10.0)
    logE = math.log10(math.e); #// for debug output
    log2pi = math.log(2.0 * math.pi)
    log2 = math.log(2.0)

    #//double logNl = logNlIn * ln10;  // Convert to base e

    #// Parition functions passed in are 2-element vectore with remperature-dependent base 10 log Us
    #// Convert to natural logs:
    #double thisLogUw, Ttheta;
    thisLogUw = 0.0 # //default initialization
    #logUw = [ 0.0 for i in range(5) ]
    logE10 = math.log(10.0)
    #print("log10UwStage ", log10UwStage)
    #for kk in range(len(logUw)):
    #    logUw[kk] = logE10*log10UwStage[kk] #// lburns new loop

    logGwL = math.log(gwL)

    #//System.out.println("chiL before: " + chiL);
    #// If we need to subtract chiI from chiL, do so *before* converting to tiny numbers in ergs!
    #////For testing with Ca II lines using gS3 internal line list only:
    #//boolean ionized = true;
    #//if (ionized) {
    #//    //System.out.println("ionized, doing chiL - chiI: " + ionized);
    #//    //         chiL = chiL - chiI;
    #//             chiL = chiL - 6.113;
    #//          }
    #//   //

    #//Log of line-center wavelength in cm
    logLam0 = math.log(lam0In) #// * 1.0e-7);

    #// energy of b-b transition
    logTransE = logH + logC - logLam0 #//ergs

    if (chiL <= 0.0):
        chiL = 1.0e-49
    logChiL = math.log(chiL) + Useful.logEv() #// Convert lower E-level from eV to ergs

    logBoltzFacL = logChiL - Useful.logK() #// Pre-factor for exponent of excitation Boltzmann factor
    boltzFacL = math.exp(logBoltzFacL)

    boltzFacGround = 0.0 / k #//I know - its zero, but let's do it this way anyway'

    #// return a 1D numDeps array of logarithmic number densities
    #// level population of lower level of bb transition (could be in either stage I or II!)

    logNums = [ 0.0 for i in range(numDeps)]

    #double num, logNum, expFac;
    #JB#
    #print("thisLogUw:",numpy.shape(logUw))
    logUwFit = ToolBox.cubicFit(masterTemp,logUw)#u(T) fit
    uw.append(logUwFit)
    #JB#

    for id in range(numDeps):

    #//Determine temperature dependenet partition functions Uw:

        #Ttheta = 5040.0 / temp[0][id]
#//NEW Determine temperature dependent partition functions Uw: lburns
        thisTemp = temp[0][id]

        """
        if (Ttheta >= 1.0):
            thisLogUw = logUw[0]

        if (Ttheta <= 0.5):
            thisLogUw = logUw[1]

        if (Ttheta > 0.5 and Ttheta < 1.0):
            thisLogUw = ( logUw[1] * (Ttheta - 0.5)/(1.0 - 0.5) ) \
                      + ( logUw[0] * (1.0 - Ttheta)/(1.0 - 0.5) )
        """
    #JB#
        thisLogUw = ToolBox.valueFromFit(logUwFit,thisTemp)#u(T) value extrapolated
    #JB#

        if (thisTemp >= 10000.0):
            thisLogUw = logUw[4]

        if (thisTemp <= 130.0):
            thisLogUw = logUw[0]

        """

        if (thisTemp > 130 and thisTemp <= 500):
            thisLogUw = logUw[1] * (thisTemp - 130)/(500 - 130) \
                      + logUw[0] * (500 - thisTemp)/(500 - 130)

        if (thisTemp > 500 and thisTemp <= 3000):
            thisLogUw = logUw[2] * (thisTemp - 500)/(3000 - 500) \
                      + logUw[1] * (3000 - thisTemp)/(3000 - 500)

        if (thisTemp > 3000 and thisTemp <= 8000):
            thisLogUw = logUw[3] * (thisTemp - 3000)/(8000 - 3000) \
                      + logUw[2] * (8000 - thisTemp)/(8000 - 3000)

        if (thisTemp > 8000 and thisTemp < 10000):
            thisLogUw = logUw[4] * (thisTemp - 8000)/(10000 - 8000) \
                      + logUw[3] * (10000 - thisTemp)/(10000 - 8000)

        """
        #print("logUw ", logUw, " thisLogUw ", thisLogUw)

        #//System.out.println("LevPops: ionized branch taken, ionized =  " + ionized);
        #// Take stat weight of ground state as partition function:
        logNums[id] = logNStage[id] - boltzFacL / temp[0][id] + logGwL - thisLogUw #// lower level of b-b transition
        #print("LevelPopsServer.stagePops id ", id, " logNStage[id] ", logNStage[id], " boltzFacL ", boltzFacL, " temp[0][id] ", temp[0][id], " logGwL ", logGwL, " thisLogUw ", thisLogUw, " logNums[id] ", logNums[id]);

        #// System.out.println("LevelPops: id, logNums[0][id], logNums[1][id], logNums[2][id], logNums[3][id]: " + id + " "
        #//          + Math.exp(logNums[0][id]) + " "
        #//         + Math.exp(logNums[1][id]) + " "
        #//          + Math.exp(logNums[2][id]) + " "
        #//        + Math.exp(logNums[3][id]));
        #//System.out.println("LevelPops: id, logNums[0][id], logNums[1][id], logNums[2][id], logNums[3][id], logNums[4][id]: " + id + " "
        #//        + logE * (logNums[0][id]) + " "
        #//        + logE * (logNums[1][id]) + " "
        #//        + logE * (logNums[2][id]) + " "
        # //        + logE * (logNums[3][id]) + " "
        #//        + logE * (logNums[4][id]) );
        #//System.out.println("LevelPops: id, logIonFracI, logIonFracII: " + id + " " + logE*logIonFracI + " " + logE*logIonFracII
        #//        + "logNum, logNumI, logNums[0][id], logNums[1][id] "
        #//        + logE*logNum + " " + logE*logNumI + " " + logE*logNums[0][id] + " " + logE*logNums[1][id]);
        #//System.out.println("LevelPops: id, logIonFracI: " + id + " " + logE*logIonFracI
        #//        + "logNums[0][id], boltzFacL/temp[0][id], logNums[2][id]: "
        #//        + logNums[0][id] + " " + boltzFacL/temp[0][id] + " " + logNums[2][id]);
    #//id loop
    #stop
    #print (uw)
    return logNums

#//This version - ionization equilibrium *WITHOUT* molecules - logNum is TOTAL element population
#def stagePops2(logNum, Ne, chiIArr, log10UwAArr,  \
#               numMols, logNumB, dissEArr, log10UwBArr, logQwABArr, logMuABArr, \
#               numDeps, temp):
def stagePops(logNum, Ne, chiIArr, logUw,  \
               numDeps, temp):
    #line 1: //species A data - ionization equilibrium of A
    #line 2: //data for set of species "B" - molecular equlibrium for set {AB}
    """Ionization equilibrium routine WITHOUT molecule formation:
    // Returns depth distribution of ionization stage populations

    // Input parameters:
    // logNum - array with depth-dependent total element number densities (cm^-3)
    // chiI1 - ground state ionization energy of neutral stage
    // chiI2 - ground state ionization energy of singly ionized stage
    // Also needs atsmopheric structure information:
    // numDeps
    // temp structure
    // rho structure
    // Atomic element A is the one whose ionization fractions are being computed
    //
    """

    ln10 = math.log(10.0)
    logE = math.log10(math.e) #// for debug output
    log2pi = math.log(2.0 * math.pi)
    log2 = math.log(2.0)

    numStages = len(chiIArr)  #// + 1; //need one more stage above the highest stage to be populated

#//    var numMols = dissEArr.length;

#// Parition functions passed in are 2-element vectore with remperature-dependent base 10 log Us
#// Convert to natural logs:
        #double Ttheta, thisTemp;
#//Default initializations:
#//We need one more stage in size of saha factor than number of stages we're actualy populating
    thisLogUw = [ 0.0 for i in range(numStages+1) ]

    for i in range(numStages+1):
        thisLogUw[i] = 0.0

    logE10 = math.log(10.0)

#//atomic ionization stage Boltzmann factors:
        #double logChiI, logBoltzFacI;
    boltzFacI = [ 0.0 for i in range(numStages) ]
    #print("numStages ", numStages, " Useful.logEv ", Useful.logEv())
    for i in range(numStages):
        #print("i ", i, " chiIArr ", chiIArr[i])
        logChiI = math.log(chiIArr[i]) + Useful.logEv()
        logBoltzFacI = logChiI  - Useful.logK()
        boltzFacI[i] = math.exp(logBoltzFacI)

    logSahaFac = log2 + (3.0 / 2.0) * (log2pi + Useful.logMe() + Useful.logK() - 2.0 * Useful.logH())

    #// return a 2D 5 x numDeps array of logarithmic number densities
    #// Row 0: neutral stage ground state population
    #// Row 1: singly ionized stage ground state population
    #// Row 2: doubly ionized stage ground state population
    #// Row 3: triply ionized stage ground state population
    #// Row 4: quadruply ionized stage ground state population
        #double[][] logNums = new double[numStages][numDeps];
    logNums = [ [ 0.0 for i in range(numDeps)] for j in range(numStages) ]

    #//We need one more stage in size of saha factor than number of stages we're actualy populating
    #//   for index accounting pirposes
    #//   For atomic ionization stages:

    logSaha = [ [ 0.0 for i in range(numStages+1)] for j in range(numStages+1) ]
    saha = [ [ 0.0 for i in range(numStages+1)] for j in range(numStages+1) ]
#//

    logIonFrac = [ 0.0 for i in range(numStages) ]
        #double expFac, logNe;

#// Now - molecular variables:

    thisLogUwA = 0.0 #// element A
    #thisLogQwAB = math.log(300.0)

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:
    logUwA = [ 0.0 for i in range(5) ]

    #JB#
    uua=[]
    #uub=[]
    #qwab=[]
    for iStg in range(numStages):
        currentUwArr=list(logUw[iStg])#u(T) determined values
        UwFit = ToolBox.cubicFit(masterTemp,currentUwArr)#u(T) fit
        uua.append(UwFit)
        #print(logUw[iStg])

    for id in range(numDeps):

        #//// reduce or enhance number density by over-all Rosseland opcity scale parameter
        #//
        #//Row 1 of Ne is log_e Ne in cm^-3
        logNe = Ne[1][id]

        #//Determine temperature dependent partition functions Uw:
        thisTemp = temp[0][id]
        #Ttheta = 5040.0 / thisTemp

    #JB#
    #use temps and partition values to create a function
    #then use said function to extrapolate values for all points
        thisLogUw[numStages] = 0.0
        for iStg in range(numStages):
            thisLogUw[iStg] = ToolBox.valueFromFit(uua[iStg],thisTemp)#u(T) value extrapolated

    #JB#

#// NEW Determine temperature dependent partition functions Uw: lburns
        if (thisTemp <= 130.0):
            for iStg in range(numStages):
                thisLogUw[iStg] = logUw[iStg][0]
            #for iMol in range(numMols):
            #    thisLogUwB[iMol] = logUwB[iMol][0]

        if (thisTemp >= 10000.0):
            for iStg in range(numStages):
                thisLogUw[iStg] = logUw[iStg][4]
            #for iMol in range(numMols):
            #    thisLogUwB[iMol] = logUwB[iMol][4]

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:
        thisLogUwA = thisLogUw[0];

        #//Ionization stage Saha factors:
        for iStg in range(numStages):
            #print("iStg ", iStg)
            logSaha[iStg+1][iStg] = logSahaFac - logNe - (boltzFacI[iStg] /temp[0][id]) + (3.0 * temp[1][id] / 2.0) + thisLogUw[iStg+1] - thisLogUw[iStg]
            saha[iStg+1][iStg] = math.exp(logSaha[iStg+1][iStg])

#//Compute log of denominator is ionization fraction, f_stage
        denominator = 1.0 #//default initialization - leading term is always unity
#//ion stage contributions:
        for jStg in range(1, numStages+1):
            addend = 1.0 #//default initialization for product series
            for iStg in range(jStg):
                #//console.log("jStg " + jStg + " saha[][] indices " + (iStg+1) + " " + iStg);
                addend = addend * saha[iStg+1][iStg]

            denominator = denominator + addend

#//
        logDenominator = math.log(denominator)

        logIonFrac[0] = -1.0 * logDenominator     #// log ionization fraction in stage I

        for jStg in range(1, numStages):
            addend = 0.0 #//default initialization for product series
            for iStg in range(jStg):
                #//console.log("jStg " + jStg + " saha[][] indices " + (iStg+1) + " " + iStg);
                addend = addend + logSaha[iStg+1][iStg]

            logIonFrac[jStg] = addend - logDenominator

        for iStg in range(numStages):
            logNums[iStg][id] = logNum[id] + logIonFrac[iStg]

    #//id loop

    return logNums;
    #//end method stagePops

#end method levelPops

#def stagePops2(logNum, Ne, chiIArr, log10UwAArr,  \
#               numMols, logNumB, dissEArr, log10UwBArr, logQwABArr, logMuABArr, \
#               numDeps, temp):
def stagePops2(logNum, Ne, chiIArr, logUw,  \
               numMols, logNumB, dissEArr, logUwB, logQwABArr, logMuABArr, \
               numDeps, temp):
    #line 1: //species A data - ionization equilibrium of A
    #line 2: //data for set of species "B" - molecular equlibrium for set {AB}
    """Ionization equilibrium routine that accounts for molecule formation:
    // Returns depth distribution of ionization stage populations

    // Input parameters:
    // logNum - array with depth-dependent total element number densities (cm^-3)
    // chiI1 - ground state ionization energy of neutral stage
    // chiI2 - ground state ionization energy of singly ionized stage
    // Also needs atsmopheric structure information:
    // numDeps
    // temp structure
    // rho structure
    // Atomic element A is the one whose ionization fractions are being computed
    //  Element B refers to array of other species with which A forms molecules AB """

    ln10 = math.log(10.0)
    logE = math.log10(math.e) #// for debug output
    log2pi = math.log(2.0 * math.pi)
    log2 = math.log(2.0)

    numStages = len(chiIArr)  #// + 1; //need one more stage above the highest stage to be populated

#//    var numMols = dissEArr.length;

#// Parition functions passed in are 2-element vectore with remperature-dependent base 10 log Us
#// Convert to natural logs:
        #double Ttheta, thisTemp;
#//Default initializations:
#//We need one more stage in size of saha factor than number of stages we're actualy populating
    thisLogUw = [ 0.0 for i in range(numStages+1) ]

    for i in range(numStages+1):
        thisLogUw[i] = 0.0

    logE10 = math.log(10.0)

#//atomic ionization stage Boltzmann factors:
        #double logChiI, logBoltzFacI;
    boltzFacI = [ 0.0 for i in range(numStages) ]
    #print("numStages ", numStages, " Useful.logEv ", Useful.logEv())
    for i in range(numStages):
        #print("i ", i, " chiIArr ", chiIArr[i])
        logChiI = math.log(chiIArr[i]) + Useful.logEv()
        logBoltzFacI = logChiI  - Useful.logK()
        boltzFacI[i] = math.exp(logBoltzFacI)

    logSahaFac = log2 + (3.0 / 2.0) * (log2pi + Useful.logMe() + Useful.logK() - 2.0 * Useful.logH())

    #// return a 2D 5 x numDeps array of logarithmic number densities
    #// Row 0: neutral stage ground state population
    #// Row 1: singly ionized stage ground state population
    #// Row 2: doubly ionized stage ground state population
    #// Row 3: triply ionized stage ground state population
    #// Row 4: quadruply ionized stage ground state population
        #double[][] logNums = new double[numStages][numDeps];
    logNums = [ [ 0.0 for i in range(numDeps)] for j in range(numStages) ]

    #//We need one more stage in size of saha factor than number of stages we're actualy populating
    #//   for index accounting pirposes
    #//   For atomic ionization stages:

    logSaha = [ [ 0.0 for i in range(numStages+1)] for j in range(numStages+1) ]
    saha = [ [ 0.0 for i in range(numStages+1)] for j in range(numStages+1) ]
#//

    logIonFrac = [ 0.0 for i in range(numStages) ]
        #double expFac, logNe;

#// Now - molecular variables:

#//Treat at least one molecule - if there are really no molecules for an atomic species,
#//there will be one phantom molecule in the denominator of the ionization fraction
#//with an impossibly high dissociation energy
    ifMols = True
    if (numMols == 0):
        ifMols = False
        numMols = 1
#//This should be inherited, but let's make sure:
        dissEArr[0] = 19.0 #//eV

#//Molecular partition functions - default initialization:
       #double[] thisLogUwB = new double[numMols];

    thisLogUwB = [ 0.0 for i in range(numMols) ]

    for iMol in range(numMols):
        thisLogUwB[iMol] = 0.0 #// variable for temp-dependent computed partn fn of array element B

    thisLogUwA = 0.0 #// element A
    thisLogQwAB = math.log(300.0)

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:
    logUwA = [ 0.0 for i in range(5) ]
    if (numMols > 0):
        for kk in range(len(logUwA)):
            logUwA[kk] = logUw[0][kk]
         #// lburns

      #//}
#//// Molecular partition functions:

#//Molecular dissociation Boltzmann factors:
    boltzFacIAB = [ 0.0 for i in range(numMols) ]
    logMolSahaFac = [ 0.0 for i in range(numMols) ]
    #//if (numMols > 0){
    #double logDissE, logBoltzFacIAB;
    for iMol in range(numMols):
        logDissE = math.log(dissEArr[iMol]) + Useful.logEv()
        logBoltzFacIAB = logDissE  - Useful.logK()
        boltzFacIAB[iMol] = math.exp(logBoltzFacIAB)
        logMolSahaFac[iMol] = (3.0 / 2.0) * (log2pi + logMuABArr[iMol] + Useful.logK() - 2.0 * Useful.logH())
        #//console.log("iMol " + iMol + " dissEArr[iMol] " + dissEArr[iMol] + " logDissE " + logE*logDissE + " logBoltzFacIAB " + logE*logBoltzFacIAB + " boltzFacIAB[iMol] " + boltzFacIAB[iMol] + " logMuABArr " + logE*logMuABArr[iMol] + " logMolSahaFac " + logE*logMolSahaFac[iMol]);

       #//}
#//   For molecular species:
    logSahaMol = [ 0.0 for i in range(numMols) ]
    invSahaMol = [ 0.0 for i in range(numMols) ]

    #JB#
    uua=[]
    uub=[]
    qwab=[]
    for iStg in range(numStages):
        currentUwArr=list(logUw[iStg])#u(T) determined values
        UwFit = ToolBox.cubicFit(masterTemp,currentUwArr)#u(T) fit
        uua.append(UwFit)
        #print(logUw[iStg])

    for iMol in range(numMols):
        currentUwBArr=list(logUwB[iMol])#u(T) determined values
        UwBFit = ToolBox.cubicFit(masterTemp,currentUwBArr)#u(T) fit
        uub.append(UwBFit)

    for id in range(numDeps):

        #//// reduce or enhance number density by over-all Rosseland opcity scale parameter
        #//
        #//Row 1 of Ne is log_e Ne in cm^-3
        logNe = Ne[1][id]

        #//Determine temperature dependent partition functions Uw:
        thisTemp = temp[0][id]
        #Ttheta = 5040.0 / thisTemp

    #JB#
    #use temps and partition values to create a function
    #then use said function to extrapolate values for all points
        thisLogUw[numStages] = 0.0
        for iStg in range(numStages):
            thisLogUw[iStg] = ToolBox.valueFromFit(uua[iStg],thisTemp)#u(T) value extrapolated
        for iMol in range(numMols):
            thisLogUwB[iMol] = ToolBox.valueFromFit(uub[iMol],thisTemp)#u(T) value extrapolated
    #JB#

#// NEW Determine temperature dependent partition functions Uw: lburns
        if (thisTemp <= 130.0):
            for iStg in range(numStages):
                thisLogUw[iStg] = logUw[iStg][0]
            for iMol in range(numMols):
                thisLogUwB[iMol] = logUwB[iMol][0]

        if (thisTemp >= 10000.0):
            for iStg in range(numStages):
                thisLogUw[iStg] = logUw[iStg][4]
            for iMol in range(numMols):
                thisLogUwB[iMol] = logUwB[iMol][4]

        for iMol in range(numMols):
            if (thisTemp < 3000.0):
                thisLogQwAB = ( logQwABArr[iMol][1] * (3000.0 - thisTemp)/(3000.0 - 500.0) ) \
                            + ( logQwABArr[iMol][2] * (thisTemp - 500.0)/(3000.0 - 500.0) )

            if ( (thisTemp >= 3000.0) and (thisTemp <= 8000.0) ):
                thisLogQwAB = ( logQwABArr[iMol][2] * (8000.0 - thisTemp)/(8000.0 - 3000.0) ) \
                            + ( logQwABArr[iMol][3] * (thisTemp - 3000.0)/(8000.0 - 3000.0) )

            if ( thisTemp > 8000.0 ):
                thisLogQwAB = ( logQwABArr[iMol][3] * (10000.0 - thisTemp)/(10000.0 - 8000.0) ) \
                            + ( logQwABArr[iMol][4] * (thisTemp - 8000.0)/(10000.0 - 8000.0) )

        #// iMol loop

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:
        thisLogUwA = thisLogUw[0];

        #//Ionization stage Saha factors:
        for iStg in range(numStages):
            #print("iStg ", iStg)
            logSaha[iStg+1][iStg] = logSahaFac - logNe - (boltzFacI[iStg] /temp[0][id]) + (3.0 * temp[1][id] / 2.0) + thisLogUw[iStg+1] - thisLogUw[iStg]
            saha[iStg+1][iStg] = math.exp(logSaha[iStg+1][iStg])

#//Molecular Saha factors:
        for iMol in range(numMols):
            logSahaMol[iMol] = logMolSahaFac[iMol] - logNumB[iMol][id] - (boltzFacIAB[iMol] / temp[0][id]) + (3.0 * temp[1][id] / 2.0) + thisLogUwB[iMol] + thisLogUwA - thisLogQwAB
#//For denominator of ionization fraction, we need *inverse* molecular Saha factors (N_AB/NI):
            logSahaMol[iMol] = -1.0 * logSahaMol[iMol]
            invSahaMol[iMol] = math.exp(logSahaMol[iMol])

#//Compute log of denominator is ionization fraction, f_stage
        denominator = 1.0 #//default initialization - leading term is always unity
#//ion stage contributions:
        for jStg in range(1, numStages+1):
            addend = 1.0 #//default initialization for product series
            for iStg in range(jStg):
                #//console.log("jStg " + jStg + " saha[][] indices " + (iStg+1) + " " + iStg);
                addend = addend * saha[iStg+1][iStg]

            denominator = denominator + addend

#//molecular contribution
        if (ifMols == True):
            for iMol in range(numMols):
                denominator = denominator + invSahaMol[iMol]

#//
        logDenominator = math.log(denominator)

        logIonFrac[0] = -1.0 * logDenominator     #// log ionization fraction in stage I

        for jStg in range(1, numStages):
            addend = 0.0 #//default initialization for product series
            for iStg in range(jStg):
                #//console.log("jStg " + jStg + " saha[][] indices " + (iStg+1) + " " + iStg);
                addend = addend + logSaha[iStg+1][iStg]

            logIonFrac[jStg] = addend - logDenominator

        for iStg in range(numStages):
            logNums[iStg][id] = logNum[id] + logIonFrac[iStg]

    #//id loop

    return logNums;
    #//end method stagePops

def stagePops3(logNum, Ne, chiIArr, logUw, numDeps, temp):
    #Version for ChromaStarPyGas: logNum is now *neutral stage* population from Phil
    # Bennett's GAS package

    #line 1: //species A data - ionization equilibrium of A
    #line 2: //data for set of species "B" - molecular equlibrium for set {AB}
    """Ionization equilibrium routine that accounts for molecule formation:
    // Returns depth distribution of ionization stage populations

    // Input parameters:
    // logNum - array with depth-dependent neutral stage  number densities (cm^-3)
    // chiI1 - ground state ionization energy of neutral stage
    // chiI2 - ground state ionization energy of singly ionized stage
    // Also needs atsmopheric structure information:
    // numDeps
    // temp structure
    // rho structure
    // Atomic element A is the one whose ionization fractions are being computed
    //  Element B refers to array of other species with which A forms molecules AB """

    ln10 = math.log(10.0)
    logE = math.log10(math.e) #// for debug output
    log2pi = math.log(2.0 * math.pi)
    log2 = math.log(2.0)

    numStages = len(chiIArr)  #// + 1; //need one more stage above the highest stage to be populated

#//    var numMols = dissEArr.length;

#// Parition functions passed in are 2-element vectore with remperature-dependent base 10 log Us
#// Convert to natural logs:
        #double Ttheta, thisTemp;
#//Default initializations:
#//We need one more stage in size of saha factor than number of stages we're actualy populating
    thisLogUw = [ 0.0 for i in range(numStages+1) ]

    for i in range(numStages+1):
        thisLogUw[i] = 0.0

    logE10 = math.log(10.0)

#//atomic ionization stage Boltzmann factors:
        #double logChiI, logBoltzFacI;
    boltzFacI = [ 0.0 for i in range(numStages) ]
    #print("numStages ", numStages, " Useful.logEv ", Useful.logEv())
    for i in range(numStages):
        #print("i ", i, " chiIArr ", chiIArr[i])
        logChiI = math.log(chiIArr[i]) + Useful.logEv()
        logBoltzFacI = logChiI  - Useful.logK()
        boltzFacI[i] = math.exp(logBoltzFacI)

    logSahaFac = log2 + (3.0 / 2.0) * (log2pi + Useful.logMe() + Useful.logK() - 2.0 * Useful.logH())

    #// return a 2D 5 x numDeps array of logarithmic number densities
    #// Row 0: neutral stage ground state population
    #// Row 1: singly ionized stage ground state population
    #// Row 2: doubly ionized stage ground state population
    #// Row 3: triply ionized stage ground state population
    #// Row 4: quadruply ionized stage ground state population
        #double[][] logNums = new double[numStages][numDeps];
    logNums = [ [ 0.0 for i in range(numDeps)] for j in range(numStages) ]

    #//We need one more stage in size of saha factor than number of stages we're actualy populating
    #//   for index accounting pirposes
    #//   For atomic ionization stages:

    #logSaha = [ [ 0.0 for i in range(numStages+1)] for j in range(numStages+1) ]
    #saha = [ [ 0.0 for i in range(numStages+1)] for j in range(numStages+1) ]
#//

    #logIonFrac = [ 0.0 for i in range(numStages) ]
        #double expFac, logNe;

    #JB#
    uua=[]
    uub=[]
    qwab=[]
    for iStg in range(numStages):
        currentUwArr=list(logUw[iStg])#u(T) determined values
        UwFit = ToolBox.cubicFit(masterTemp,currentUwArr)#u(T) fit
        uua.append(UwFit)
        #print(logUw[iStg])

    for id in range(numDeps):

        #//// reduce or enhance number density by over-all Rosseland opcity scale parameter
        #//
        #//Row 1 of Ne is log_e Ne in cm^-3
        logNe = Ne[1][id]

        #//Determine temperature dependent partition functions Uw:
        thisTemp = temp[0][id]
        #Ttheta = 5040.0 / thisTemp

    #JB#
    #use temps and partition values to create a function
    #then use said function to extrapolate values for all points
        thisLogUw[numStages] = 0.0
        for iStg in range(numStages):
            thisLogUw[iStg] = ToolBox.valueFromFit(uua[iStg],thisTemp)#u(T) value extrapolated
    #JB#

#// NEW Determine temperature dependent partition functions Uw: lburns
        if (thisTemp <= 130.0):
            for iStg in range(numStages):
                thisLogUw[iStg] = logUw[iStg][0]

        if (thisTemp >= 10000.0):
            for iStg in range(numStages):
                thisLogUw[iStg] = logUw[iStg][4]

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:
        #thisLogUwA = thisLogUw[0];

        #//Ionization stage Saha factors:
        logNums[0][id] = logNum[id]
        for iStg in range(1, numStages):
            #print("iStg ", iStg)
            thisLogSaha = logSahaFac - logNe - (boltzFacI[iStg-1] /temp[0][id]) + (3.0 * temp[1][id] / 2.0) + thisLogUw[iStg] - thisLogUw[iStg-1]
            #saha[iStg+1][iStg] = math.exp(logSaha[iStg+1][iStg])
            logNums[iStg][id] = logNums[iStg-1][id] + thisLogSaha

    #//id loop

    return logNums;
    #//end method stagePops

#def sahaRHS(chiI, log10UwUArr, log10UwLArr, temp):
def sahaRHS(chiI, logUwU, logUwL, temp):

    """RHS of partial pressure formulation of Saha equation in standard form (N_U*P_e/N_L on LHS)
 // Returns depth distribution of LHS: Phi(T) === N_U*P_e/N_L (David Gray notation)

// Input parameters:
// chiI - ground state ionization energy of lower stage
// log10UwUArr, log10UwLArr - array of temperature-dependent partition function for upper and lower ionization stage
// Also needs atsmopheric structure information:
// numDeps
// temp structure
//
// Atomic element "A" is the one whose ionization fractions are being computed
//  Element "B" refers to array of other species with which A forms molecules "AB" """

    ln10 = math.log(10.0)
    logE = math.log10(math.e) #// for debug output
    log2pi = math.log(2.0 * math.pi)
    log2 = math.log(2.0)

#//    var numMols = dissEArr.length;

#// Parition functions passed in are 2-element vectore with remperature-dependent base 10 log Us
#// Convert to natural logs:
    #double Ttheta, thisTemp;
#//Default initializations:
#//We need one more stage in size of saha factor than number of stages we're actualy populating
    thisLogUwU = 0.0
    thisLogUwL = 0.0

    logE10 = math.log(10.0)
#//We need one more stage in size of saha factor than number of stages we're actualy populating
    #logUwU = [0.0 for i in range(5)]
    #logUwL = [0.0 for i in range(5)]
    for kk in range(len(logUwL)):
        logUwU[kk] = logUwL[kk]
    #   logUwL[kk] = logE10*log10UwLArr[kk]

    #//System.out.println("chiL before: " + chiL);
    #// If we need to subtract chiI from chiL, do so *before* converting to tiny numbers in ergs!

#//atomic ionization stage Boltzmann factors:
    #double logChiI, logBoltzFacI;
    #double boltzFacI;
    logChiI = math.log(chiI) + Useful.logEv()
    logBoltzFacI = logChiI  - Useful.logK()
    boltzFacI = math.exp(logBoltzFacI)

#//Extra factor of k to get k^5/2 in the P_e formulation of Saha Eq.
    logSahaFac = log2 + (3.0 / 2.0) * (log2pi + Useful.logMe() + Useful.logK() - 2.0 * Useful.logH()) + Useful.logK()

    #//double[] logLHS = new double[numDeps];
    #double logLHS;

#//   For atomic ionization stages:
    #double logSaha, saha, expFac;

#//  for (int id = 0; id < numDeps; id++) {

#//
#//Determine temperature dependent partition functions Uw:
    thisTemp = temp[0]
    #Ttheta = 5040.0 / thisTemp

    """
    if (Ttheta >= 1.0):
        thisLogUwU = logUwU[0]
        thisLogUwL = logUwL[0]

    if (Ttheta <= 0.5):
        thisLogUwU = logUwU[1]
        thisLogUwL = logUwL[1]

    if (Ttheta > 0.5 and Ttheta < 1.0):
        thisLogUwU = ( logUwU[1] * (Ttheta - 0.5)/(1.0 - 0.5) )
        + ( logUwU[0] * (1.0 - Ttheta)/(1.0 - 0.5) )
        thisLogUwL = ( logUwL[1] * (Ttheta - 0.5)/(1.0 - 0.5) )
        + ( logUwL[0] * (1.0 - Ttheta)/(1.0 - 0.5) )
    """
    #JB#
    currentUwUArr=list(logUwU)#u(T) determined values
    UwUFit = ToolBox.cubicFit(masterTemp,currentUwUArr)#u(T) fit
    thisLogUwU = ToolBox.valueFromFit(UwUFit,thisTemp)#u(T) value extrapolated

    currentUwLArr=list(logUwL)#u(T) determined values
    UwLFit = ToolBox.cubicFit(masterTemp,currentUwLArr)#u(T) fit
    thisLogUwL = ToolBox.valueFromFit(UwLFit,thisTemp)#u(T) value extrapolated
    #JB#

    #will need to do this one in Main as it goes through its own loop of temp
    #if thisTemp == superTemp[0][len(superTemp[0])]:
    #    uwu.append(UwUFit)
    #    uwl.append(UwLFit)
    #
    #JB#

    if (thisTemp <= 130.0):
        thisLogUwU = logUwU[0]
        thisLogUwL = logUwL[0]

    if (thisTemp >= 10000.0):
        thisLogUwU = logUwU[4]
        thisLogUwL = logUwL[4]

    """
    if (thisTemp > 130 and thisTemp <= 500):
        thisLogUwU = logUwU[1] * (thisTemp - 130)/(500 - 130) \
            + logUwU[0] * (500 - thisTemp)/(500 - 130)
        thisLogUwL = logUwL[1] * (thisTemp - 130)/(500 - 130) \
            + logUwL[0] * (500 - thisTemp)/(500 - 130)

    if (thisTemp > 500 and thisTemp <= 3000):
        thisLogUwU = logUwU[2] * (thisTemp - 500)/(3000 - 500) \
            + logUwU[1] * (3000 - thisTemp)/(3000 - 500)
        thisLogUwL = logUwL[2] * (thisTemp - 500)/(3000 - 500) \
            + logUwL[1] * (3000 - thisTemp)/(3000 - 500)

    if (thisTemp > 3000 and thisTemp <= 8000):
        thisLogUwU = logUwU[3] * (thisTemp - 3000)/(8000 - 3000) \
            + logUwU[2] * (8000 - thisTemp)/(8000 - 3000)
        thisLogUwL = logUwL[3] * (thisTemp - 3000)/(8000 - 3000) \
            + logUwL[2] * (8000 - thisTemp)/(8000 - 3000)

    if (thisTemp > 8000 and thisTemp < 10000):
        thisLogUwU = logUwU[4] * (thisTemp - 8000)/(10000 - 8000) \
            + logUwU[3] * (10000 - thisTemp)/(10000 - 8000)
        thisLogUwL = logUwL[4] * (thisTemp - 8000)/(10000 - 8000) \
            + logUwL[3] * (10000 - thisTemp)/(10000 - 8000)

    if (thisTemp >= 10000):
        thisLogUwU = logUwU[4]
        thisLogUwL = logUwL[4]
    """

#//Ionization stage Saha factors:

#//Need T_kin^5/2 in the P_e formulation of Saha Eq.
    logSaha = logSahaFac - (boltzFacI /temp[0]) + (5.0 * temp[1] / 2.0) + thisLogUwU - thisLogUwL
    #// saha = Math.exp(logSaha);

    #//logLHS[id] = logSaha;
    logLHS = logSaha;
#//    } //id loop

    return logLHS;
    #JB
    #return [logLHS,[[UwUFit,thisLogUwU],[UwLFit,thisLogUwL]]]
#//
#    } //end method sahaRHS

#def molPops(nmrtrLogNumB, nmrtrDissE, log10UwA, nmrtrLog10UwB, nmrtrLogQwAB, nmrtrLogMuAB, \
#            numMolsB, logNumB, dissEArr, log10UwBArr, logQwABArr, logMuABArr,   \
#            logGroundRatio, numDeps, temp):
def molPops(nmrtrLogNumB, nmrtrDissE, logUwA, nmrtrLogUwB, nmrtrLogQwAB, nmrtrLogMuAB, \
            numMolsB, logNumB, dissEArr, logUwB, logQwABArr, logMuABArr,   \
            logGroundRatio, numDeps, temp):
    # line 1: //species A data - ionization equilibrium of A
    # //data for set of species "B" - molecular equlibrium for set {AB}

    """Diatomic molecular equilibrium routine that accounts for molecule formation:
 // Returns depth distribution of molecular population

// Input parameters:
// logNum - array with depth-dependent total element number densities (cm^-3)
// chiI1 - ground state ionization energy of neutral stage
// chiI2 - ground state ionization energy of singly ionized stage
// Also needs atsmopheric structure information:
// numDeps
// temp structure
// rho structure
//
// Atomic element "A" is the one kept on the LHS of the master fraction, whose ionization fractions are included
//   in the denominator of the master fraction
//  Element "B" refers to array of other sintpecies with which A forms molecules "AB" """

    logE = math.log10(math.e) #// for debug output
    #//System.out.println("molPops: nmrtrDissE " + nmrtrDissE + " log10UwA " + log10UwA[0] + " " + log10UwA[1] + " nmrtrLog10UwB " +
    #//     nmrtrLog10UwB[0] + " " + nmrtrLog10UwB[1] + " nmrtrLog10QwAB " + logE*nmrtrLogQwAB[2] + " nmrtrLogMuAB " + logE*nmrtrLogMuAB
    #//     + " numMolsB " + numMolsB + " dissEArr " + dissEArr[0] + " log10UwBArr " + log10UwBArr[0][0] + " " + log10UwBArr[0][1] + " log10QwABArr " +
    #//     logE*logQwABArr[0][2] + " logMuABArr " + logE*logMuABArr[0]);
    #//System.out.println("Line: nmrtrLog10UwB[0] " + logE*nmrtrLog10UwB[0] + " nmrtrLog10UwB[1] " + logE*nmrtrLog10UwB[1]);

    ln10 = math.log(10.0)
    log2pi = math.log(2.0 * math.pi)
    log2 = math.log(2.0)

    logE10 = math.log(10.0)
#// Convert to natural logs:
    #double Ttheta, thisTemp;

#//Treat at least one molecule - if there are really no molecules for an atomic species,
#//there will be one phantom molecule in the denominator of the ionization fraction
#//with an impossibly high dissociation energy
    if (numMolsB == 0):
        numMolsB = 1
#//This should be inherited, but let's make sure:
        dissEArr[0] = 29.0 #//eV

    #//var molPops = function(logNum, numeratorLogNumB, numeratorDissE, numeratorLog10UwA, numeratorLog10QwAB, numeratorLogMuAB,  //species A data - ionization equilibrium of A
    #//Molecular partition functions - default initialization:
    thisLogUwB = [0.0 for i in range(numMolsB)]

    for iMol in range(numMolsB):
        thisLogUwB[iMol] = 0.0 #// variable for temp-dependent computed partn fn of array element B

    thisLogUwA = 0.0 #// element A
    nmrtrThisLogUwB = 0.0 #// element A
    thisLogQwAB = math.log(300.0)
    nmrtrThisLogQwAB = math.log(300.0)

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:
    #logUwA = [0.0 for i in range(5)]

    #nmrtrLogUwB = [0.0 for i in range(5)]

    #for kk in range(len(logUwA)):
        #logUwA[kk] = logE10*log10UwA[kk]
        #nmrtrLogUwB[kk] = logE10*nmrtrLog10UwB[kk]

        #// lburns

#// Array of elements B for all molecular species AB:
    #double[][] logUwB = new double[numMolsB][2];
    #logUwB = [ [ 0.0 for i in range(5) ] for j in range(numMolsB) ]
    #//if (numMolsB > 0){
    #for iMol in range(numMolsB):
    #    for kk in range(5):
    #        logUwB[iMol][kk] = logE10*log10UwBArr[iMol][kk]
           # // lburns new loop

    #//}
#// Molecular partition functions:
#//       double nmrtrLogQwAB = logE10*nmrtrLog10QwAB;
#//       double[] logQwAB = new double[numMolsB];
#//      //if (numMolsB > 0){
#//       for (int iMol = 0; iMol < numMolsB; iMol++){
#//          logQwAB[iMol] = logE10*log10QwABArr[iMol];
#//       }
#      //}
#//Molecular dissociation Boltzmann factors:
    nmrtrBoltzFacIAB = 0.0
    nmrtrLogMolSahaFac = 0.0
    logDissE = math.log(nmrtrDissE)  + Useful.logEv()
    #//System.out.println("logDissE " + logE*logDissE)
    logBoltzFacIAB = logDissE  - Useful.logK()
    #//System.out.println("logBoltzFacIAB " + logE*logBoltzFacIAB);
    nmrtrBoltzFacIAB = math.exp(logBoltzFacIAB)
    nmrtrLogMolSahaFac = (3.0 / 2.0) * (log2pi + nmrtrLogMuAB  + Useful.logK() - 2.0 * Useful.logH())
    #//System.out.println("nmrtrLogMolSahaFac " + logE*nmrtrLogMolSahaFac);
    #//System.out.println("nmrtrDissE " + nmrtrDissE + " logDissE " + logE*logDissE + " logBoltzFacIAB " + logE*logBoltzFacIAB + " nmrtrBoltzFacIAB " + nmrtrBoltzFacIAB + " nmrtrLogMuAB " + logE*nmrtrLogMuAB + " nmrtrLogMolSahaFac " + logE*nmrtrLogMolSahaFac);

    boltzFacIAB = [0.0 for i in range(numMolsB)]
    logMolSahaFac = [0.0 for i in range(numMolsB)]
    #//if (numMolsB > 0){
    for iMol in range(numMolsB):
        logDissE = math.log(dissEArr[iMol]) + Useful.logEv()
        logBoltzFacIAB = logDissE  - Useful.logK()
        boltzFacIAB[iMol] = math.exp(logBoltzFacIAB)
        logMolSahaFac[iMol] = (3.0 / 2.0) * (log2pi + logMuABArr[iMol] + Useful.logK() - 2.0 * Useful.logH())
  #//System.out.println("logMolSahaFac[iMol] " + logE*logMolSahaFac[iMol]);
  #//System.out.println("iMol " + iMol + " dissEArr[iMol] " + dissEArr[iMol] + " logDissE " + logE*logDissE + " logBoltzFacIAB " + logE*logBoltzFacIAB + " boltzFacIAB[iMol] " + boltzFacIAB[iMol] + " logMuABArr " + logE*logMuABArr[iMol] + " logMolSahaFac " + logE*logMolSahaFac[iMol]);

  #//double[] logNums = new double[numDeps]

  #//}
#//   For molecular species:
    #double nmrtrSaha, nmrtrLogSahaMol, nmrtrLogInvSahaMol; //, nmrtrInvSahaMol;
    logMolFrac = [0.0 for i in range(numDeps)]
    logSahaMol = [0.0 for i in range(numMolsB)]
    invSahaMol = [0.0 for i in range(numMolsB)]

    #JB#
    currentUwAArr=list(logUwA)#u(T) determined values
    UwAFit = ToolBox.cubicFit(masterTemp, currentUwAArr)#u(T) fit

    nmrtrLogUwBArr=list(nmrtrLogUwB)#u(T) determined values
    nmrtrLogUwBFit = ToolBox.cubicFit(masterTemp, nmrtrLogUwBArr)#u(T) fit

    #uwa.append(UwAFit)
    #uwb.append(nmrtrLogUwBFit)

    uwbFits=[]
    qwabFit = []
    for iMol in range(numMolsB):
          currentUwBArr=list(logUwB[iMol])
          UwBFit = ToolBox.cubicFit(masterTemp, currentUwBArr)
          uwbFits.append(UwBFit)

          currentLogQwABArr=list(logQwABArr[iMol])#u(T) determined values
          QwABFit = ToolBox.cubicFit(masterTemp, currentLogQwABArr)#u(T) fit
          qwabFit.append(QwABFit)

    #nmrtrQwABArr=list(nmrtrLogQwAB)#u(T) determined values
    #nmrtrQwABFit = ToolBox.cubicFit(masterTemp, nmrtrQwABArr)#u(T) fit

    #for Mols in range(numMolsB):
    #    currentLogUwBArr=list(logUwB[Mols])#u(T) determined values
    #    UwBFit=cubicFit(masterTemp,currentLogUwBArr)#u(T) fit

    #JB#
#//

    temps=[]
    #valb=[]
    #vala=[]
    #valnb=[]
    #valqab=[]
    #valnmrtrqwb=[]
    #//  System.out.println("molPops: id      nmrtrLogNumB      logNumBArr[0]      logGroundRatio");
    for id in range(numDeps):

        #//System.out.format("%03d, %21.15f, %21.15f, %21.15f, %n", id, logE*nmrtrLogNumB[id], logE*logNumB[0][id], logE*logGroundRatio[id]);

        #//// reduce or enhance number density by over-all Rosseland opcity scale parameter

#//Determine temparature dependent partition functions Uw:
        thisTemp = temp[0][id]
        temps.append(thisTemp)
        #Ttheta = 5040.0 / thisTemp
        """
        if (Ttheta >= 1.0):
            thisLogUwA = logUwA[0]
            nmrtrThisLogUwB = nmrtrLogUwB[0]
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][0]

        if (Ttheta <= 0.5):
            thisLogUwA = logUwA[1]
            nmrtrThisLogUwB = nmrtrLogUwB[1]
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][1]

        if (Ttheta > 0.5 and Ttheta < 1.0):
            thisLogUwA = ( logUwA[1] * ((Ttheta - 0.5)/(1.0 - 0.5)) ) \
            + ( logUwA[0] * ((1.0 - Ttheta)/(1.0 - 0.5)) )
            nmrtrThisLogUwB = ( nmrtrLogUwB[1] * ((Ttheta - 0.5)/(1.0 - 0.5)) ) \
            + ( nmrtrLogUwB[0] * ((1.0 - Ttheta)/(1.0 - 0.5)) )
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = ( logUwB[iMol][1] * ((Ttheta - 0.5)/(1.0 - 0.5)) ) \
                + ( logUwB[iMol][0] * ((1.0 - Ttheta)/(1.0 - 0.5)) )
        """
    #JB#
        thisLogUwA = float(ToolBox.valueFromFit(UwAFit,thisTemp))#u(T) value extrapolated
        #vala.append(thisLogUwA)

        nmrtrThisLogUwB = float(ToolBox.valueFromFit(nmrtrLogUwBFit,thisTemp))#u(T) value extrapolated
        #valnb.append(nmrtrThisLogUwB)

    #for iMol in range(numMolsB):
    #    thisLogUwB[iMol]=logUwB[iMol]

        for iMol in range(numMolsB):
            thisLogUwB[iMol] = ToolBox.valueFromFit(uwbFits[iMol],thisTemp)#u(T) value extrapolated
            #valb.append(thisLogUwB[iMol])

#// NEW Determine temperature dependent partition functions Uw: lburns
        thisTemp = temp[0][id]
        if (thisTemp <= 130.0):
            thisLogUwA = logUwA[0]
            nmrtrThisLogUwB = nmrtrLogUwB[0]
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][0]
        if (thisTemp >= 10000.0):
            thisLogUwA = logUwA[4]
            nmrtrThisLogUwB = nmrtrLogUwB[4]
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][4]
        """
        if (thisTemp > 130 and thisTemp <= 500):
            thisLogUwA = logUwA[1] * (thisTemp - 130)/(500 - 130) \
                       + logUwA[0] * (500 - thisTemp)/(500 - 130)
            nmrtrThisLogUwB = nmrtrLogUwB[1] * (thisTemp - 130)/(500 - 130) \
                            + nmrtrLogUwB[0] * (500 - thisTemp)/(500 - 130)
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][1] * (thisTemp - 130)/(500 - 130) \
                                 + logUwB[iMol][0] * (500 - thisTemp)/(500 - 130)

        if (thisTemp > 500 and thisTemp <= 3000):
            thisLogUwA = logUwA[2] * (thisTemp - 500)/(3000 - 500) \
                       + logUwA[1] * (3000 - thisTemp)/(3000 - 500)
            nmrtrThisLogUwB = nmrtrLogUwB[2] * (thisTemp - 500)/(3000 - 500) \
                            + nmrtrLogUwB[1] * (3000 - thisTemp)/(3000 - 500)
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][2] * (thisTemp - 500)/(3000 - 500) \
                                 + logUwB[iMol][1] * (3000 - thisTemp)/(3000 - 500)

        if (thisTemp > 3000 and thisTemp <= 8000):
            thisLogUwA = logUwA[3] * (thisTemp - 3000)/(8000 - 3000) \
                       + logUwA[2] * (8000 - thisTemp)/(8000 - 3000)
            nmrtrThisLogUwB = nmrtrLogUwB[3] * (thisTemp - 3000)/(8000 - 3000) \
                            + nmrtrLogUwB[2] * (8000 - thisTemp)/(8000 - 3000)
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][3] * (thisTemp - 3000)/(8000 - 3000) \
                                 + logUwB[iMol][2] * (8000 - thisTemp)/(8000 - 3000)

        if (thisTemp > 8000 and thisTemp < 10000):
            thisLogUwA = logUwA[4] * (thisTemp - 8000)/(10000 - 8000) \
                       + logUwA[3] * (10000 - thisTemp)/(10000 - 8000)
            nmrtrThisLogUwB = nmrtrLogUwB[4] * (thisTemp - 8000)/(10000 - 8000) \
                            + nmrtrLogUwB[3] * (10000 - thisTemp)/(10000 - 8000)
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][4] * (thisTemp - 8000)/(10000 - 8000) \
                                + logUwB[iMol][3] * (10000 - thisTemp)/(10000 - 8000)

        if (thisTemp >= 10000):
            thisLogUwA = logUwA[4]
            nmrtrThisLogUwB = nmrtrLogUwB[4]
            for iMol in range(numMolsB):
                thisLogUwB[iMol] = logUwB[iMol][4]

        """

        #iMol loops for Q's

        for iMol in range(numMolsB):
            if (thisTemp < 3000.0):
                thisLogQwAB = ( logQwABArr[iMol][1] * (3000.0 - thisTemp)/(3000.0 - 500.0) ) \
                + ( logQwABArr[iMol][2] * (thisTemp - 500.0)/(3000.0 - 500.0) )
            if ( (thisTemp >= 3000.0) and (thisTemp <= 8000.0) ):
                thisLogQwAB = ( logQwABArr[iMol][2] * (8000.0 - thisTemp)/(8000.0 - 3000.0) ) \
                + ( logQwABArr[iMol][3] * (thisTemp - 3000.0)/(8000.0 - 3000.0) )
            if ( thisTemp > 8000.0 ):
                thisLogQwAB = ( logQwABArr[iMol][3] * (10000.0 - thisTemp)/(10000.0 - 8000.0) ) \
                + ( logQwABArr[iMol][4] * (thisTemp - 8000.0)/(10000.0 - 8000.0) )

            if (thisTemp < 3000.0):
                nmrtrThisLogQwAB = ( nmrtrLogQwAB[1] * (3000.0 - thisTemp)/(3000.0 - 500.0) ) \
                + ( nmrtrLogQwAB[2] * (thisTemp - 500.0)/(3000.0 - 500.0) )
            if ( (thisTemp >= 3000.0) and (thisTemp <= 8000.0) ):
                nmrtrThisLogQwAB = ( nmrtrLogQwAB[2] * (8000.0 - thisTemp)/(8000.0 - 3000.0) ) \
                + ( nmrtrLogQwAB[3] * (thisTemp - 3000.0)/(8000.0 - 3000.0) )
            if ( thisTemp > 8000.0 ):
                nmrtrThisLogQwAB = ( nmrtrLogQwAB[3] * (10000.0 - thisTemp)/(10000.0 - 8000.0) ) \
                + ( nmrtrLogQwAB[4] * (thisTemp - 8000.0)/(10000.0 - 8000.0) )

#//For clarity: neutral stage of atom whose ionization equilibrium is being computed is element A
#// for molecule formation:

#   //Ionization stage Saha factors:
#//System.out.println("id " + id + " nmrtrLogNumB[id] " + logE*nmrtrLogNumB[id]);
#  // if (id == 16){
#  //   System.out.println("id " + id + " nmrtrLogNumB[id] " + logE*nmrtrLogNumB[id] + " pp nmrtB " + (logE*(nmrtrLogNumB[id]+temp[1][id]+Useful.logK())) + " nmrtrThisLogUwB " + logE*nmrtrThisLogUwB + " thisLogUwA " + logE*thisLogUwA + " nmrtrLogQwAB " + logE*nmrtrThisLogQwAB);
#   //System.out.println("nmrtrThisLogUwB " + logE*nmrtrThisLogUwB + " thisLogUwA " + logE*thisLogUwA + " nmrtrThisLogQwAB " + logE*nmrtrThisLogQwAB);
#   // }
        nmrtrLogSahaMol = nmrtrLogMolSahaFac - nmrtrLogNumB[id] - (nmrtrBoltzFacIAB / temp[0][id]) + (3.0 * temp[1][id] / 2.0) + nmrtrThisLogUwB + thisLogUwA - nmrtrThisLogQwAB
        nmrtrLogInvSahaMol = -1.0 * nmrtrLogSahaMol
        #//System.out.println("nmrtrLogInvSahaMol " + logE*nmrtrLogInvSahaMol);
        #//nmrtrInvSahaMol = Math.exp(nmrtrLogSahaMol);
        #//   if (id == 16){
        #//       System.out.println("nmrtrLogInvSahaMol " + logE*nmrtrLogInvSahaMol);
        #//   }
        #//   if (id == 16){
        #//        System.out.println("nmrtrBoltzFacIAB " + nmrtrBoltzFacIAB + " nmrtrThisLogUwB " + logE*nmrtrThisLogUwB + " thisLogUwA " + logE*thisLogUwA + " nmrtrThisLogQwAB " + nmrtrThisLogQwAB);
        #//        System.out.println("nmrtrLogSahaMol " + logE*nmrtrLogSahaMol); // + " nmrtrInvSahaMol " + nmrtrInvSahaMol);
        #//   }

#//Molecular Saha factors:
        for iMol in range(numMolsB):
#//System.out.println("iMol " + iMol + " id " + id + " logNumB[iMol][id] " + logE*nmrtrLogNumB[id]);
#//System.out.println("iMol " + iMol + " thisLogUwB[iMol] " + logE*thisLogUwB[iMol] + " thisLogUwA " + logE*thisLogUwA + " thisLogQwAB " + logE*thisLogQwAB);
            logSahaMol[iMol] = logMolSahaFac[iMol] - logNumB[iMol][id] - (boltzFacIAB[iMol] / temp[0][id]) + (3.0 * temp[1][id] / 2.0) + float(thisLogUwB[iMol]) + thisLogUwA - thisLogQwAB
#//For denominator of ionization fraction, we need *inverse* molecular Saha factors (N_AB/NI):
            logSahaMol[iMol] = -1.0 * logSahaMol[iMol]
            invSahaMol[iMol] = math.exp(logSahaMol[iMol])
            #//TEST invSahaMol[iMol] = 1.0e-99; //test
     #//     if (id == 16){
     #//         System.out.println("iMol " + iMol + " boltzFacIAB[iMol] " + boltzFacIAB[iMol] + " thisLogUwB[iMol] " + logE*thisLogUwB[iMol] + " logQwAB[iMol] " + logE*thisLogQwAB + " logNumB[iMol][id] " + logE*logNumB[iMol][id] + " logMolSahaFac[iMol] " + logE*logMolSahaFac[iMol]);
     #//         System.out.println("iMol " + iMol + " logSahaMol " + logE*logSahaMol[iMol] + " invSahaMol[iMol] " + invSahaMol[iMol]);
     #//     }

#//Compute log of denominator is ionization fraction, f_stage
#        //default initialization
#        //  - ratio of total atomic particles in all ionization stages to number in ground state:
        denominator = math.exp(logGroundRatio[id]) #//default initialization - ratio of total atomic particles in all ionization stages to number in ground state
#//molecular contribution
        for iMol in range(numMolsB):
            #//  if (id == 16){
            #//   System.out.println("invSahaMol[iMol] " + invSahaMol[iMol] + " denominator " + denominator);
            #//  }
            denominator = denominator + invSahaMol[iMol]

#//
        logDenominator = math.log(denominator)
        #//System.out.println("logGroundRatio[id] " + logE*logGroundRatio[id] + " logDenominator " + logE*logDenominator);
        #//  if (id == 16){
        #// System.out.println("id " + id + " logGroundRatio " + logGroundRatio[id] + " logDenominator " + logDenominator);
        #//  }
        #//if (id == 36){
        #//     System.out.println("logDenominator " + logE*logDenominator);
        #// }
        #//var logDenominator = Math.log( 1.0 + saha21 + (saha32 * saha21) + (saha43 * saha32 * saha21) + (saha54 * saha43 * saha32 * saha21) );

        logMolFrac[id] = nmrtrLogInvSahaMol - logDenominator
        #//  if (id == 16){
        #//       System.out.println("id " + id + " logMolFrac[id] " + logE*logMolFrac[id]);
        #//  }

        #//logNums[id] = logNum[id] + logMolFrac;
    #} //id loop

    #JB - check (never used)#
    #print(uwa)
    #print(uwb)
    #title("logUwA")
    """
    plot(temps,vala)
    tempT=[]
    for t in masterTemp:
        tempT.append(valueFromFit(UwAFit,t))
    scatter(masterTemp,(tempT))
    show()

    #title("nmrtrlogUwB")
    plot(temps,valnb)
    tempT=[]
    for t in masterTemp:
        tempT.append(valueFromFit(nmrtrLogUwBFit,t))
    scatter(masterTemp,(tempT))
    show()

    #title("logUwB")
    plot(temps,valb)
    tempT=[]
    for t in masterTemp:
        tempT.append(valueFromFit(UwBFit,t))
    scatter(masterTemp,(tempT))
    show()

    #title("logQwAB")
    plot(temps,valqab)
    tempT=[]
    for t in masterTemp:
        tempT.append(valueFromFit(QwABFit,t))
    scatter(masterTemp,(tempT))
    show()

    #title("nmrtrlogQwAB")
    plot(temps,valnmrtrqwb)
    tempT=[]
    for t in masterTemp:
        tempT.append(valueFromFit(nmrtrQwABFit,t))
    scatter(masterTemp,(tempT))
    show()
    """
    #JB#
    return logMolFrac
    #//end method stagePops


#  Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
"""This showcases how simple it is to build image classification networks.

It follows description from this TensorFlow tutorial:
    https://www.tensorflow.org/versions/master/tutorials/mnist/pros/index.html#deep-mnist-for-experts
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from sklearn import metrics
import tensorflow as tf

layers = tf.contrib.layers
learn = tf.contrib.learn

def max_pool_2x2(tensor_in):
  return tf.nn.max_pool(
      tensor_in, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')

def conv_model(feature, target, mode):
  """2-layer convolution model."""
  # Convert the target to a one-hot tensor of shape (batch_size, 10) and
  # with a on-value of 1 for each one-hot vector of length 10.
  target = tf.one_hot(tf.cast(target, tf.int32), 10, 1, 0)

  # Reshape feature to 4d tensor with 2nd and 3rd dimensions being
  # image width and height final dimension being the number of color channels.
  feature = tf.reshape(feature, [-1, 28, 28, 1])

  # First conv layer will compute 32 features for each 5x5 patch
  with tf.variable_scope('conv_layer1'):
    h_conv1 = layers.convolution2d(
        feature, 32, kernel_size=[5, 5], activation_fn=tf.nn.relu)
    h_pool1 = max_pool_2x2(h_conv1)

  # Second conv layer will compute 64 features for each 5x5 patch.
  with tf.variable_scope('conv_layer2'):
    h_conv2 = layers.convolution2d(
        h_pool1, 64, kernel_size=[5, 5], activation_fn=tf.nn.relu)
    h_pool2 = max_pool_2x2(h_conv2)
    # reshape tensor into a batch of vectors
    h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])

  # Densely connected layer with 1024 neurons.
  h_fc1 = layers.dropout(
      layers.fully_connected(
          h_pool2_flat, 1024, activation_fn=tf.nn.relu),
      keep_prob=0.5,
      is_training=mode == tf.contrib.learn.ModeKeys.TRAIN)

  # Compute logits (1 per class) and compute loss.
  logits = layers.fully_connected(h_fc1, 10, activation_fn=None)
  loss = tf.losses.softmax_cross_entropy(target, logits)

  # Create a tensor for training op.
  train_op = layers.optimize_loss(
      loss,
      tf.contrib.framework.get_global_step(),
      optimizer='SGD',
      learning_rate=0.001)

  return tf.argmax(logits, 1), loss, train_op

def main(unused_args):
  ### Download and load MNIST dataset.
  mnist = learn.datasets.load_dataset('mnist')

  ### Linear classifier.
  feature_columns = learn.infer_real_valued_columns_from_input(
      mnist.train.images)
  classifier = learn.LinearClassifier(
      feature_columns=feature_columns, n_classes=10)
  classifier.fit(mnist.train.images,
                 mnist.train.labels.astype(np.int32),
                 batch_size=100,
                 steps=1000)
  score = metrics.accuracy_score(mnist.test.labels,
                                 list(classifier.predict(mnist.test.images)))
  print('Accuracy: {0:f}'.format(score))

  ### Convolutional network
  classifier = learn.Estimator(model_fn=conv_model)
  classifier.fit(mnist.train.images,
                 mnist.train.labels,
                 batch_size=100,
                 steps=20000)
  score = metrics.accuracy_score(mnist.test.labels,
                                 list(classifier.predict(mnist.test.images)))
  print('Accuracy: {0:f}'.format(score))

if __name__ == '__main__':
  tf.app.run()


import csv
import os
import random
import numpy as np
import pandas as pd
import matlab_wrapper
from lib.utils.timer import Timer
from tools.ear_recog import get_gt, ROI_boxes
import scipy.io as sio

def listdir_no_hidden(path):
    list1 = []
    for f in sorted(os.listdir(path)):
        if not f.startswith('.'):
            p = os.path.abspath(path)
            list1.append(os.path.join(p, f))
    return list1

def write_list_to_csv(list1, path_out, header=False):
    temp = pd.DataFrame(list1)
    temp.to_csv(path_out, index=False, header=header)

def save_gt_roidb_csv(data_path, csv_path, image_index_output_path, gt_output_path, test_image_path, test_gt):
    box_list = pd.read_csv(csv_path, header=0).get_values()

    image_path_list = listdir_no_hidden(data_path)
    assert len(box_list) == len(image_path_list), 'the length of box list must equal to image list'
    new_list = []
    new_list1 = []
    for idx, entry in enumerate(image_path_list):
        s1 = str(entry)
        temp = box_list[idx]
        # change the x y coordination to correct [X1 Y1 X2 Y2]
        x1 = str(temp[-2])
        y1 = str(temp[-4])
        x2 = str(temp[-1])
        y2 = str(temp[-3])
        s2 = x1+' '+ y1+' '+x2+' '+y2
        new_list.append(s1 + ' 1 ' + s2)
        new_list1.append(s1)
    # shuffle the idx of training set
    shuffle_idx = range(len(image_path_list))
    random.seed(641) # make it can be reproduce
    random.shuffle(shuffle_idx)
    train_idx = shuffle_idx[0:437]
    test_idx = shuffle_idx[437:]

    train_image_path = [new_list1[idx] for idx in train_idx]
    train_gt = [new_list[idx] for idx in train_idx]
    test_image_path_data = [new_list1[idx] for idx in test_idx]
    test_gt_data = [new_list[idx] for idx in test_idx]

    write_list_to_csv(train_gt, gt_output_path)
    write_list_to_csv(train_image_path, image_index_output_path)

    write_list_to_csv(test_gt_data, test_gt)
    write_list_to_csv(test_image_path_data, test_image_path)
def initialize_matlab():
    matlab = matlab_wrapper.MatlabSession()

    # edge_detector OP_method
    matlab.eval("cd('/home/harrysocool/Github/fast-rcnn/OP_methods/edges')")
    matlab.eval("addpath(genpath('/home/harrysocool/Github/fast-rcnn/OP_methods/edges'))")
    matlab.eval("toolboxCompile")

    # # selective_search OP_method
    # matlab.eval("cd('/home/harrysocool/Github/fast-rcnn/OP_methods/selective_search_ijcv_with_python')")
    # matlab.eval("addpath(genpath('/home/harrysocool/Github/fast-rcnn/OP_methods/selective_search_ijcv_with_python'))")

    return matlab

def time_analyse(matlab, cmd, image_filepath, par1, par2):
    timer = Timer()
    timer.tic()

    obj_proposals = ROI_boxes(matlab, image_filepath, cmd, par1, par2)

    timer.toc()
    time = timer.total_time
    box_numer = len(obj_proposals)

    return time, box_numer, obj_proposals

def mean_IOU_ratio(image_index, dets):
    ratio = np.empty(0,dtype=np.float64)
    (x1, y1, x2, y2) = get_gt(image_index)
    if dets.size > 4:
        for box in dets:
            X1 = box[0]
            Y1 = box[1]
            X2 = box[2]
            Y2 = box[3]
            if ((np.float32(x1)-X1)<=15 and (X2- np.float32(x2))<=15
                and (np.float32(y1)-Y1)<=15 and (Y2-np.float32(y2))<=15):
                ratio = np.append(ratio,1.0)
            else:
                SI = max(0, min(x2, X2) - max(x1, X1)) * \
                     max(0, min(y2, Y2) - max(y1, Y1))
                SU = (x2 - x1) * (y2 - y1) + (X2 - X1) * (Y2 - Y1) - SI
                ratio = np.append(ratio, SI/SU)
    if ratio.size == 0:
        big_ratio = 0
    else:
        big = np.where(ratio >= 0.1)[0].size
        total = float(len(dets))
        big_ratio = float(big/total)
    return big_ratio

if __name__ == '__main__':
    datasets_path = '/home/harrysocool/Github/fast-rcnn/DatabaseEars'
    csv_path = os.path.join(datasets_path, 'boundaries.csv')
    image_path = os.path.join(datasets_path, 'DatabaseEars/')
    gt_output_path = os.path.join(datasets_path, '../','ear_recognition/data_file/gt_roidb.csv')
    image_index_output_path = os.path.join(datasets_path, '../', 'ear_recognition/data_file/image_index_list.csv')
    mat_output_filename = os.path.join(datasets_path, '../','ear_recognition/data_file/all_boxes.mat')

    test_gt_output_path = os.path.join(datasets_path, '../','ear_recognition/data_file/test_gt_roidb.csv')
    test_image_index_output_path = os.path.join(datasets_path, '../', 'ear_recognition/data_file/test_image_index_list.csv')

    # save_gt_roidb_csv(image_path, csv_path, image_index_output_path, gt_output_path, test_image_index_output_path,
    #                   test_gt_output_path)

    matlab = initialize_matlab()
    timer = Timer()

    list1 = pd.read_csv(test_image_index_output_path, header=None).values.flatten().tolist()
    cmd = 'ss'
    # ks = [50 100 150 200 300];
    par2_list = [8]
    # par2_list = [3]
    time_csv_out_path = os.path.join(os.path.dirname(datasets_path), 'result', cmd + '_' + 'OPtune_result_1.csv')
    if not os.path.exists(time_csv_out_path):
        write_list_to_csv(par2_list, time_csv_out_path)
    with open(time_csv_out_path, 'a') as csvfile:
        writer = csv.writer(csvfile)
        list2 = []
        for par2 in [7]:
            for par1 in [7]:
                # par1 = float(par1)/100
                # all_boxes = np.zeros((437,), dtype=np.object)
                for index, image_path in enumerate(list1):
                    # if index>300:
                    #     break
                    time, box_numer, obj_proposals = time_analyse(matlab, cmd, image_path, par1, par2)
                    ratio = mean_IOU_ratio(index + 1, obj_proposals)
                    # list2.append([time, box_numer])
                    # print('{} has processed in {:.3f} seconds with {} boxes'.format(len(list2), time, box_numer))
                    print('No. {} has processed with par {} {}, box {} IOU ratio {:.3f} in {:.2f} seconds'.format(index,
                                                                                par1, par2,box_numer ,ratio, time))
                    writer.writerow([par1, par2,ratio,box_numer, time])
                    # all_boxes[index] = obj_proposals

    # sio.savemat(mat_output_filename, {'all_boxes': all_boxes})

    # write_list_to_csv(list2, time_csv_out_path)

    # fnames_cell = "{" + ",".join("'{}'".format(x) for x in list1) + "}"
    # command = "res = {}({}, '{}')".format('selective_search', fnames_cell, mat_output_filename)
    # print(command)
    # #
    # matlab.eval(command)


# -*- coding: utf-8 -*-

import numpy as np
import pandas.lib as lib
import pandas.util.testing as tm

from pandas.compat import long, u

class TestParseSQL(tm.TestCase):

    def test_convert_sql_column_floats(self):
        arr = np.array([1.5, None, 3, 4.2], dtype=object)
        result = lib.convert_sql_column(arr)
        expected = np.array([1.5, np.nan, 3, 4.2], dtype='f8')
        self.assert_numpy_array_equal(result, expected)

    def test_convert_sql_column_strings(self):
        arr = np.array(['1.5', None, '3', '4.2'], dtype=object)
        result = lib.convert_sql_column(arr)
        expected = np.array(['1.5', np.nan, '3', '4.2'], dtype=object)
        self.assert_numpy_array_equal(result, expected)

    def test_convert_sql_column_unicode(self):
        arr = np.array([u('1.5'), None, u('3'), u('4.2')],
                       dtype=object)
        result = lib.convert_sql_column(arr)
        expected = np.array([u('1.5'), np.nan, u('3'), u('4.2')],
                            dtype=object)
        self.assert_numpy_array_equal(result, expected)

    def test_convert_sql_column_ints(self):
        arr = np.array([1, 2, 3, 4], dtype='O')
        arr2 = np.array([1, 2, 3, 4], dtype='i4').astype('O')
        result = lib.convert_sql_column(arr)
        result2 = lib.convert_sql_column(arr2)
        expected = np.array([1, 2, 3, 4], dtype='i8')
        self.assert_numpy_array_equal(result, expected)
        self.assert_numpy_array_equal(result2, expected)

        arr = np.array([1, 2, 3, None, 4], dtype='O')
        result = lib.convert_sql_column(arr)
        expected = np.array([1, 2, 3, np.nan, 4], dtype='f8')
        self.assert_numpy_array_equal(result, expected)

    def test_convert_sql_column_longs(self):
        arr = np.array([long(1), long(2), long(3), long(4)], dtype='O')
        result = lib.convert_sql_column(arr)
        expected = np.array([1, 2, 3, 4], dtype='i8')
        self.assert_numpy_array_equal(result, expected)

        arr = np.array([long(1), long(2), long(3), None, long(4)], dtype='O')
        result = lib.convert_sql_column(arr)
        expected = np.array([1, 2, 3, np.nan, 4], dtype='f8')
        self.assert_numpy_array_equal(result, expected)

    def test_convert_sql_column_bools(self):
        arr = np.array([True, False, True, False], dtype='O')
        result = lib.convert_sql_column(arr)
        expected = np.array([True, False, True, False], dtype=bool)
        self.assert_numpy_array_equal(result, expected)

        arr = np.array([True, False, None, False], dtype='O')
        result = lib.convert_sql_column(arr)
        expected = np.array([True, False, np.nan, False], dtype=object)
        self.assert_numpy_array_equal(result, expected)

    def test_convert_sql_column_decimals(self):
        from decimal import Decimal
        arr = np.array([Decimal('1.5'), None, Decimal('3'), Decimal('4.2')])
        result = lib.convert_sql_column(arr)
        expected = np.array([1.5, np.nan, 3, 4.2], dtype='f8')
        self.assert_numpy_array_equal(result, expected)

    def test_convert_downcast_int64(self):
        from pandas.parser import na_values

        arr = np.array([1, 2, 7, 8, 10], dtype=np.int64)
        expected = np.array([1, 2, 7, 8, 10], dtype=np.int8)

        # default argument
        result = lib.downcast_int64(arr, na_values)
        self.assert_numpy_array_equal(result, expected)

        result = lib.downcast_int64(arr, na_values, use_unsigned=False)
        self.assert_numpy_array_equal(result, expected)

        expected = np.array([1, 2, 7, 8, 10], dtype=np.uint8)
        result = lib.downcast_int64(arr, na_values, use_unsigned=True)
        self.assert_numpy_array_equal(result, expected)

        # still cast to int8 despite use_unsigned=True
        # because of the negative number as an element
        arr = np.array([1, 2, -7, 8, 10], dtype=np.int64)
        expected = np.array([1, 2, -7, 8, 10], dtype=np.int8)
        result = lib.downcast_int64(arr, na_values, use_unsigned=True)
        self.assert_numpy_array_equal(result, expected)

        arr = np.array([1, 2, 7, 8, 300], dtype=np.int64)
        expected = np.array([1, 2, 7, 8, 300], dtype=np.int16)
        result = lib.downcast_int64(arr, na_values)
        self.assert_numpy_array_equal(result, expected)

        int8_na = na_values[np.int8]
        int64_na = na_values[np.int64]
        arr = np.array([int64_na, 2, 3, 10, 15], dtype=np.int64)
        expected = np.array([int8_na, 2, 3, 10, 15], dtype=np.int8)
        result = lib.downcast_int64(arr, na_values)
        self.assert_numpy_array_equal(result, expected)

if __name__ == '__main__':
    import nose

    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)


__author__ = 'IVMIT KFU: Gataullin Ravil & Veselovkiy Sergei'

from copy import copy
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from  time import time
import warnings
warnings.filterwarnings("ignore")
from sklearn.cross_validation import train_test_split
from structure import Position

class PatchVarianceClassifier:
    def __init__(self, init_patch):
        self.init_patch_variance = np.var(init_patch.content)

    def classify(self, patch):
        # return 1 if object is positive detected
        # return 0 if object is negative detected
        if np.var(patch.content) > 0.5 * self.init_patch_variance:
            return 1
        else:
            return 0

    def predict_patch(self, patch):
        return np.var(patch.content) / self.init_patch_variance

    def predict_position(self, position):
        return np.var(position.calculate_patch().content) / self.init_patch_variance

class EnsembleClassifier:
    def __init__(self, learning_component):
        self.learning_component = learning_component
        self.classifier = RandomForestClassifier(max_depth=3)

    def classify(self, patch):
        # return 1 if object is positive detected
        # return 0 if object is negative detected
        feature = patch.calculate_feature(self.learning_component.descriptor)
        if self.classifier.predict_proba(feature)[0][self.positive_class_index] > 0.5:
            return 1
        else:
            return 0

    def predict_patch(self, patch):
        feature = patch.calculate_feature(self.learning_component.descriptor)
        return self.classifier.predict_proba(feature)[0][self.positive_class_index]

    def predict_position(self, position):
        feature = position.calculate_patch().calculate_feature(self.learning_component.descriptor)
        return self.classifier.predict_proba(feature)[0][self.positive_class_index]

    def relearn(self, test_size=0):
        samples, weights, targets = self.learning_component.get_training_set(const_weight=True)
        train_samples, test_samples, train_targets, test_targets = train_test_split(samples, targets, test_size=test_size, random_state=np.random.RandomState(0))
        count_positives = 1.0*np.count_nonzero(train_targets)
        count_negatives = 1.0*(len(train_targets) - count_positives)
        positive_weight = count_negatives/len(train_targets)
        negative_weight = count_positives/len(train_targets)
        weights = np.array([positive_weight if target == 1 else negative_weight for target in train_targets])
        self.classifier.fit(train_samples, train_targets, sample_weight=weights)
        self.learning_component.new_samples_count = 0
        if len(test_samples) > 0:
            test_result = [self.classifier.predict(sample) for sample in test_samples]
            true_positives = 0.0
            count_test_positives = 1.0*np.count_nonzero(test_targets)
            count_result_positives = 1.0*np.count_nonzero(test_result)
            for i in xrange(len(test_targets)):
                if test_targets[i] == test_result[i] and test_result[i] == 1:
                    true_positives += 1
            precision = true_positives / count_test_positives
            recall = true_positives / count_result_positives
            print "Precision:", precision
            print "Recall", recall
            if precision + recall != 0:
                print "F-score:", 2 * precision * recall / (precision + recall)
            else:
                print "F-score:", 0
        self.positive_class_index = 0
        for elem in self.classifier.classes_:
            if elem != 1.0:
                self.positive_class_index += 1
            else:
                break

class NearestNeighborClassifier:

    def __init__(self, learning_component, lmbd = 0.1, tetta = 0.6):
        self.learning_component = learning_component
        self.lmbd = lmbd
        self.tetta = tetta

    def classify(self, patch):
        # return 1 if object is positive detected
        # return 0 if object is negative detected
        if self.learning_component.relative_similarity(patch) > self.tetta:
            return 1
        else:
            return 0

    def predict_patch(self, patch):
        return self.learning_component.relative_similarity(patch)

    def predict_position(self, position):
        return self.learning_component.relative_similarity(position.calculate_patch())

def scanning_window(init_position, scales_step = 1.2, slip_step = 0.1, minimal_bounding_box_size = 20, min_step=1, max_step=20):
    flag_inc = True
    flag_dec = False
    position = copy(init_position)
    while min(position.width, position.height) >= minimal_bounding_box_size:
        position.update(x=0,y=0)
        step_width = min(max(min_step,int(slip_step * position.width)),max_step)
        step_height = min(max(min_step,int(slip_step * position.height)),max_step)
        while position.is_correct():
            while position.is_correct():
                yield position
                position.update(x=position.x+step_width)
            position.update(x=0, y=position.y+step_height)
        # if position.is_correct():
        #     yield position
        # is_end = False
        # step_width = int(slip_step * position.width)
        # step_height = int(slip_step * position.height)
        # layer = 1
        # xx = position.x
        # yy = position.y
        # while not is_end:
        #     is_end = True
        #     for start_point, vector in (([-1,-1],[1,0]),([1,-1],[0,1]),([1,1],[-1,0]),([-1,1],[0,-1])):
        #         position.update(x=xx + (start_point[0]*layer + vector[0])*step_width, y=yy+(start_point[1]*layer + vector[1])*step_height)
        #         while position.is_correct() and xx - layer*step_width <= position.x <= xx + layer*step_width and yy - layer*step_height <= position.y <= yy + layer*step_height:
        #             is_end = False
        #             yield position
        #             position.update(x=position.x+vector[0]*step_width, y=position.y+vector[1]*step_height)
        #     layer += 1
        if flag_inc:
            position.update(height=int(position.height * scales_step), width = int(position.width * scales_step))
        if position.height > position.buffer[0].shape[0] or position.width > position.buffer[0].shape[0]:
            flag_inc = False
            flag_dec = True
            position = copy(init_position)
        if flag_dec:
            position.update(height=int(position.height / scales_step), width = int(position.width / scales_step))

def get_sliding_positions(init_position, scales_step = 1.2, slip_step = 0.1, minimal_bounding_box_size = 20, min_step=2, max_step=2):
    sliding_positions = []
    flag_inc = True
    flag_dec = False
    position = copy(init_position)
    while min(position.width, position.height) >= minimal_bounding_box_size:
        position.update(x=0,y=0)
        step_width = min(max(min_step,int(slip_step * position.width)),max_step)
        step_height = min(max(min_step,int(slip_step * position.height)),max_step)
        while position.is_correct():
            while position.is_correct():
                sliding_positions.append(copy(position))
                position.update(x=position.x+step_width)
            position.update(x=0, y=position.y+step_height)
        if flag_inc:
            position.update(height=int(position.height * scales_step), width = int(position.width * scales_step))
        if position.height > position.buffer[0].shape[0] or position.width > position.buffer[0].shape[0]:
            flag_inc = False
            flag_dec = True
            position = copy(init_position)
        if flag_dec:
            position.update(height=int(position.height / scales_step), width = int(position.width / scales_step))
    return sliding_positions

class Detector:
    def __init__(self, init_position, learning_component, threshold_patch_variance=0.5, threshold_ensemble=0.5, threshold_nearest_neighbor=0.6):
        self.learning_component = learning_component
        self.patch_variance_classifier = PatchVarianceClassifier(learning_component.init_patch)
        self.ensemble_classifier = EnsembleClassifier(learning_component)
        self.nearest_neighbor_classifier = NearestNeighborClassifier(learning_component)
        self.threshold_patch_variance = threshold_patch_variance
        self.threshold_ensemble = threshold_ensemble
        self.threshold_nearest_neighbor = threshold_nearest_neighbor
        self.sliding_positions = get_sliding_positions(init_position, scales_step = 1.2, slip_step = 0.1, minimal_bounding_box_size = 50, min_step=2, max_step=10)

    def cascaded_classifier(self, patch):
        # 3 stages of classify
        # return 1 if object is positive detected
        # return 0 if object is negative detected
        if self.patch_variance_classifier.predict_patch(patch) < self.threshold_patch_variance:
            return 0
        if self.ensemble_classifier.predict_patch(patch) < self.threshold_patch_variance:
            return 0
        # elif self.nearest_neighbor_classifier.predict_patch(patch) < self.threshold_nearest_neighbor:
        #     return 0
        return 1

    def detect(self, position, is_tracked):
        if self.learning_component.new_samples_count > 10:
            start = time()
            self.ensemble_classifier.relearn()
            print "Relearn:", time() - start
        detected_windows = []
        predict_times = []
        for current_position in self.sliding_positions:
            start = time()
            proba = self.predict_position(current_position)
            predict_times.append(time() - start)
            if proba == 1:
                detected_windows.append((current_position.get_window(), current_position.calculate_patch(), proba))
                self.learning_component.add_new_positive(current_position.calculate_patch())
                if is_tracked:
                    return detected_windows
            else:
                self.learning_component.add_new_negative(current_position.calculate_patch())
        print "Analysed window count:", len(predict_times)
        print "Max detection time:", np.max(predict_times)
        print "Min detection time:", np.min(predict_times)
        print "Mean detection time:", np.mean(predict_times)
        return detected_windows

    def predict_patch(self, patch):
        return self.cascaded_classifier(patch)

    def predict_position(self, position):
        return self.cascaded_classifier(position.calculate_patch())


from nose.tools import assert_equal

import numpy as np
from scipy import linalg

from sklearn.cross_validation import train_test_split
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_less
from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import ignore_warnings
from sklearn.utils.testing import assert_no_warnings, assert_warns
from sklearn.utils.testing import TempMemmap
from sklearn.utils import ConvergenceWarning
from sklearn import linear_model, datasets
from sklearn.linear_model.least_angle import _lars_path_residues

diabetes = datasets.load_diabetes()
X, y = diabetes.data, diabetes.target

# TODO: use another dataset that has multiple drops

def test_simple():
    # Principle of Lars is to keep covariances tied and decreasing

    # also test verbose output
    from sklearn.externals.six.moves import cStringIO as StringIO
    import sys
    old_stdout = sys.stdout
    try:
        sys.stdout = StringIO()

        alphas_, active, coef_path_ = linear_model.lars_path(
            diabetes.data, diabetes.target, method="lar", verbose=10)

        sys.stdout = old_stdout

        for (i, coef_) in enumerate(coef_path_.T):
            res = y - np.dot(X, coef_)
            cov = np.dot(X.T, res)
            C = np.max(abs(cov))
            eps = 1e-3
            ocur = len(cov[C - eps < abs(cov)])
            if i < X.shape[1]:
                assert_true(ocur == i + 1)
            else:
                # no more than max_pred variables can go into the active set
                assert_true(ocur == X.shape[1])
    finally:
        sys.stdout = old_stdout

def test_simple_precomputed():
    # The same, with precomputed Gram matrix

    G = np.dot(diabetes.data.T, diabetes.data)
    alphas_, active, coef_path_ = linear_model.lars_path(
        diabetes.data, diabetes.target, Gram=G, method="lar")

    for i, coef_ in enumerate(coef_path_.T):
        res = y - np.dot(X, coef_)
        cov = np.dot(X.T, res)
        C = np.max(abs(cov))
        eps = 1e-3
        ocur = len(cov[C - eps < abs(cov)])
        if i < X.shape[1]:
            assert_true(ocur == i + 1)
        else:
            # no more than max_pred variables can go into the active set
            assert_true(ocur == X.shape[1])

def test_all_precomputed():
    # Test that lars_path with precomputed Gram and Xy gives the right answer
    X, y = diabetes.data, diabetes.target
    G = np.dot(X.T, X)
    Xy = np.dot(X.T, y)
    for method in 'lar', 'lasso':
        output = linear_model.lars_path(X, y, method=method)
        output_pre = linear_model.lars_path(X, y, Gram=G, Xy=Xy, method=method)
        for expected, got in zip(output, output_pre):
            assert_array_almost_equal(expected, got)

def test_lars_lstsq():
    # Test that Lars gives least square solution at the end
    # of the path
    X1 = 3 * diabetes.data  # use un-normalized dataset
    clf = linear_model.LassoLars(alpha=0.)
    clf.fit(X1, y)
    coef_lstsq = np.linalg.lstsq(X1, y)[0]
    assert_array_almost_equal(clf.coef_, coef_lstsq)

def test_lasso_gives_lstsq_solution():
    # Test that Lars Lasso gives least square solution at the end
    # of the path
    alphas_, active, coef_path_ = linear_model.lars_path(X, y, method="lasso")
    coef_lstsq = np.linalg.lstsq(X, y)[0]
    assert_array_almost_equal(coef_lstsq, coef_path_[:, -1])

def test_collinearity():
    # Check that lars_path is robust to collinearity in input
    X = np.array([[3., 3., 1.],
                  [2., 2., 0.],
                  [1., 1., 0]])
    y = np.array([1., 0., 0])

    f = ignore_warnings
    _, _, coef_path_ = f(linear_model.lars_path)(X, y, alpha_min=0.01)
    assert_true(not np.isnan(coef_path_).any())
    residual = np.dot(X, coef_path_[:, -1]) - y
    assert_less((residual ** 2).sum(), 1.)  # just make sure it's bounded

    n_samples = 10
    X = np.random.rand(n_samples, 5)
    y = np.zeros(n_samples)
    _, _, coef_path_ = linear_model.lars_path(X, y, Gram='auto', copy_X=False,
                                              copy_Gram=False, alpha_min=0.,
                                              method='lasso', verbose=0,
                                              max_iter=500)
    assert_array_almost_equal(coef_path_, np.zeros_like(coef_path_))

def test_no_path():
    # Test that the ``return_path=False`` option returns the correct output

    alphas_, active_, coef_path_ = linear_model.lars_path(
        diabetes.data, diabetes.target, method="lar")
    alpha_, active, coef = linear_model.lars_path(
        diabetes.data, diabetes.target, method="lar", return_path=False)

    assert_array_almost_equal(coef, coef_path_[:, -1])
    assert_true(alpha_ == alphas_[-1])

def test_no_path_precomputed():
    # Test that the ``return_path=False`` option with Gram remains correct

    G = np.dot(diabetes.data.T, diabetes.data)

    alphas_, active_, coef_path_ = linear_model.lars_path(
        diabetes.data, diabetes.target, method="lar", Gram=G)
    alpha_, active, coef = linear_model.lars_path(
        diabetes.data, diabetes.target, method="lar", Gram=G,
        return_path=False)

    assert_array_almost_equal(coef, coef_path_[:, -1])
    assert_true(alpha_ == alphas_[-1])

def test_no_path_all_precomputed():
    # Test that the ``return_path=False`` option with Gram and Xy remains
    # correct
    X, y = 3 * diabetes.data, diabetes.target
    G = np.dot(X.T, X)
    Xy = np.dot(X.T, y)

    alphas_, active_, coef_path_ = linear_model.lars_path(
        X, y, method="lasso", Gram=G, Xy=Xy, alpha_min=0.9)
    print("---")
    alpha_, active, coef = linear_model.lars_path(
        X, y, method="lasso", Gram=G, Xy=Xy, alpha_min=0.9, return_path=False)

    assert_array_almost_equal(coef, coef_path_[:, -1])
    assert_true(alpha_ == alphas_[-1])

def test_singular_matrix():
    # Test when input is a singular matrix
    X1 = np.array([[1, 1.], [1., 1.]])
    y1 = np.array([1, 1])
    alphas, active, coef_path = linear_model.lars_path(X1, y1)
    assert_array_almost_equal(coef_path.T, [[0, 0], [1, 0]])

def test_rank_deficient_design():
    # consistency test that checks that LARS Lasso is handling rank
    # deficient input data (with n_features < rank) in the same way
    # as coordinate descent Lasso
    y = [5, 0, 5]
    for X in ([[5, 0],
               [0, 5],
               [10, 10]],

              [[10, 10, 0],
               [1e-32, 0, 0],
               [0, 0, 1]],
              ):
        # To be able to use the coefs to compute the objective function,
        # we need to turn off normalization
        lars = linear_model.LassoLars(.1, normalize=False)
        coef_lars_ = lars.fit(X, y).coef_
        obj_lars = (1. / (2. * 3.)
                    * linalg.norm(y - np.dot(X, coef_lars_)) ** 2
                    + .1 * linalg.norm(coef_lars_, 1))
        coord_descent = linear_model.Lasso(.1, tol=1e-6, normalize=False)
        coef_cd_ = coord_descent.fit(X, y).coef_
        obj_cd = ((1. / (2. * 3.)) * linalg.norm(y - np.dot(X, coef_cd_)) ** 2
                  + .1 * linalg.norm(coef_cd_, 1))
        assert_less(obj_lars, obj_cd * (1. + 1e-8))

def test_lasso_lars_vs_lasso_cd(verbose=False):
    # Test that LassoLars and Lasso using coordinate descent give the
    # same results.
    X = 3 * diabetes.data

    alphas, _, lasso_path = linear_model.lars_path(X, y, method='lasso')
    lasso_cd = linear_model.Lasso(fit_intercept=False, tol=1e-8)
    for c, a in zip(lasso_path.T, alphas):
        if a == 0:
            continue
        lasso_cd.alpha = a
        lasso_cd.fit(X, y)
        error = linalg.norm(c - lasso_cd.coef_)
        assert_less(error, 0.01)

    # similar test, with the classifiers
    for alpha in np.linspace(1e-2, 1 - 1e-2, 20):
        clf1 = linear_model.LassoLars(alpha=alpha, normalize=False).fit(X, y)
        clf2 = linear_model.Lasso(alpha=alpha, tol=1e-8,
                                  normalize=False).fit(X, y)
        err = linalg.norm(clf1.coef_ - clf2.coef_)
        assert_less(err, 1e-3)

    # same test, with normalized data
    X = diabetes.data
    alphas, _, lasso_path = linear_model.lars_path(X, y, method='lasso')
    lasso_cd = linear_model.Lasso(fit_intercept=False, normalize=True,
                                  tol=1e-8)
    for c, a in zip(lasso_path.T, alphas):
        if a == 0:
            continue
        lasso_cd.alpha = a
        lasso_cd.fit(X, y)
        error = linalg.norm(c - lasso_cd.coef_)
        assert_less(error, 0.01)

def test_lasso_lars_vs_lasso_cd_early_stopping(verbose=False):
    # Test that LassoLars and Lasso using coordinate descent give the
    # same results when early stopping is used.
    # (test : before, in the middle, and in the last part of the path)
    alphas_min = [10, 0.9, 1e-4]
    for alphas_min in alphas_min:
        alphas, _, lasso_path = linear_model.lars_path(X, y, method='lasso',
                                                       alpha_min=0.9)
        lasso_cd = linear_model.Lasso(fit_intercept=False, tol=1e-8)
        lasso_cd.alpha = alphas[-1]
        lasso_cd.fit(X, y)
        error = linalg.norm(lasso_path[:, -1] - lasso_cd.coef_)
        assert_less(error, 0.01)

    alphas_min = [10, 0.9, 1e-4]
    # same test, with normalization
    for alphas_min in alphas_min:
        alphas, _, lasso_path = linear_model.lars_path(X, y, method='lasso',
                                                       alpha_min=0.9)
        lasso_cd = linear_model.Lasso(fit_intercept=True, normalize=True,
                                      tol=1e-8)
        lasso_cd.alpha = alphas[-1]
        lasso_cd.fit(X, y)
        error = linalg.norm(lasso_path[:, -1] - lasso_cd.coef_)
        assert_less(error, 0.01)

def test_lasso_lars_path_length():
    # Test that the path length of the LassoLars is right
    lasso = linear_model.LassoLars()
    lasso.fit(X, y)
    lasso2 = linear_model.LassoLars(alpha=lasso.alphas_[2])
    lasso2.fit(X, y)
    assert_array_almost_equal(lasso.alphas_[:3], lasso2.alphas_)
    # Also check that the sequence of alphas is always decreasing
    assert_true(np.all(np.diff(lasso.alphas_) < 0))

def test_lasso_lars_vs_lasso_cd_ill_conditioned():
    # Test lasso lars on a very ill-conditioned design, and check that
    # it does not blow up, and stays somewhat close to a solution given
    # by the coordinate descent solver
    # Also test that lasso_path (using lars_path output style) gives
    # the same result as lars_path and previous lasso output style
    # under these conditions.
    rng = np.random.RandomState(42)

    # Generate data
    n, m = 70, 100
    k = 5
    X = rng.randn(n, m)
    w = np.zeros((m, 1))
    i = np.arange(0, m)
    rng.shuffle(i)
    supp = i[:k]
    w[supp] = np.sign(rng.randn(k, 1)) * (rng.rand(k, 1) + 1)
    y = np.dot(X, w)
    sigma = 0.2
    y += sigma * rng.rand(*y.shape)
    y = y.squeeze()
    lars_alphas, _, lars_coef = linear_model.lars_path(X, y, method='lasso')

    _, lasso_coef2, _ = linear_model.lasso_path(X, y,
                                                alphas=lars_alphas,
                                                tol=1e-6,
                                                fit_intercept=False)

    assert_array_almost_equal(lars_coef, lasso_coef2, decimal=1)

def test_lasso_lars_vs_lasso_cd_ill_conditioned2():
    # Create an ill-conditioned situation in which the LARS has to go
    # far in the path to converge, and check that LARS and coordinate
    # descent give the same answers
    # Note it used to be the case that Lars had to use the drop for good
    # strategy for this but this is no longer the case with the
    # equality_tolerance checks
    X = [[1e20, 1e20, 0],
         [-1e-32, 0, 0],
         [1, 1, 1]]
    y = [10, 10, 1]
    alpha = .0001

    def objective_function(coef):
        return (1. / (2. * len(X)) * linalg.norm(y - np.dot(X, coef)) ** 2
                + alpha * linalg.norm(coef, 1))

    lars = linear_model.LassoLars(alpha=alpha, normalize=False)
    assert_warns(ConvergenceWarning, lars.fit, X, y)
    lars_coef_ = lars.coef_
    lars_obj = objective_function(lars_coef_)

    coord_descent = linear_model.Lasso(alpha=alpha, tol=1e-10, normalize=False)
    cd_coef_ = coord_descent.fit(X, y).coef_
    cd_obj = objective_function(cd_coef_)

    assert_less(lars_obj, cd_obj * (1. + 1e-8))

def test_lars_add_features():
    # assure that at least some features get added if necessary
    # test for 6d2b4c
    # Hilbert matrix
    n = 5
    H = 1. / (np.arange(1, n + 1) + np.arange(n)[:, np.newaxis])
    clf = linear_model.Lars(fit_intercept=False).fit(
        H, np.arange(n))
    assert_true(np.all(np.isfinite(clf.coef_)))

def test_lars_n_nonzero_coefs(verbose=False):
    lars = linear_model.Lars(n_nonzero_coefs=6, verbose=verbose)
    lars.fit(X, y)
    assert_equal(len(lars.coef_.nonzero()[0]), 6)
    # The path should be of length 6 + 1 in a Lars going down to 6
    # non-zero coefs
    assert_equal(len(lars.alphas_), 7)

def test_multitarget():
    # Assure that estimators receiving multidimensional y do the right thing
    X = diabetes.data
    Y = np.vstack([diabetes.target, diabetes.target ** 2]).T
    n_targets = Y.shape[1]

    for estimator in (linear_model.LassoLars(), linear_model.Lars()):
        estimator.fit(X, Y)
        Y_pred = estimator.predict(X)
        Y_dec = estimator.decision_function(X)
        assert_array_almost_equal(Y_pred, Y_dec)
        alphas, active, coef, path = (estimator.alphas_, estimator.active_,
                                      estimator.coef_, estimator.coef_path_)
        for k in range(n_targets):
            estimator.fit(X, Y[:, k])
            y_pred = estimator.predict(X)
            assert_array_almost_equal(alphas[k], estimator.alphas_)
            assert_array_almost_equal(active[k], estimator.active_)
            assert_array_almost_equal(coef[k], estimator.coef_)
            assert_array_almost_equal(path[k], estimator.coef_path_)
            assert_array_almost_equal(Y_pred[:, k], y_pred)

def test_lars_cv():
    # Test the LassoLarsCV object by checking that the optimal alpha
    # increases as the number of samples increases.
    # This property is not actually garantied in general and is just a
    # property of the given dataset, with the given steps chosen.
    old_alpha = 0
    lars_cv = linear_model.LassoLarsCV()
    for length in (400, 200, 100):
        X = diabetes.data[:length]
        y = diabetes.target[:length]
        lars_cv.fit(X, y)
        np.testing.assert_array_less(old_alpha, lars_cv.alpha_)
        old_alpha = lars_cv.alpha_

def test_lasso_lars_ic():
    # Test the LassoLarsIC object by checking that
    # - some good features are selected.
    # - alpha_bic > alpha_aic
    # - n_nonzero_bic < n_nonzero_aic
    lars_bic = linear_model.LassoLarsIC('bic')
    lars_aic = linear_model.LassoLarsIC('aic')
    rng = np.random.RandomState(42)
    X = diabetes.data
    y = diabetes.target
    X = np.c_[X, rng.randn(X.shape[0], 4)]  # add 4 bad features
    lars_bic.fit(X, y)
    lars_aic.fit(X, y)
    nonzero_bic = np.where(lars_bic.coef_)[0]
    nonzero_aic = np.where(lars_aic.coef_)[0]
    assert_greater(lars_bic.alpha_, lars_aic.alpha_)
    assert_less(len(nonzero_bic), len(nonzero_aic))
    assert_less(np.max(nonzero_bic), diabetes.data.shape[1])

    # test error on unknown IC
    lars_broken = linear_model.LassoLarsIC('<unknown>')
    assert_raises(ValueError, lars_broken.fit, X, y)

def test_no_warning_for_zero_mse():
    # LassoLarsIC should not warn for log of zero MSE.
    y = np.arange(10, dtype=float)
    X = y.reshape(-1, 1)
    lars = linear_model.LassoLarsIC(normalize=False)
    assert_no_warnings(lars.fit, X, y)
    assert_true(np.any(np.isinf(lars.criterion_)))

def test_lars_path_readonly_data():
    # When using automated memory mapping on large input, the
    # fold data is in read-only mode
    # This is a non-regression test for:
    # https://github.com/scikit-learn/scikit-learn/issues/4597
    splitted_data = train_test_split(X, y, random_state=42)
    with TempMemmap(splitted_data) as (X_train, X_test, y_train, y_test):
        # The following should not fail despite copy=False
        _lars_path_residues(X_train, y_train, X_test, y_test, copy=False)

def test_lars_path_positive_constraint():
    # this is the main test for the positive parameter on the lars_path method
    # the estimator classes just make use of this function

    # we do the test on the diabetes dataset

    # ensure that we get negative coefficients when positive=False
    # and all positive when positive=True
    # for method 'lar' (default) and lasso
    for method in ['lar', 'lasso']:
        alpha, active, coefs = \
            linear_model.lars_path(diabetes['data'], diabetes['target'],
                                   return_path=True, method=method,
                                   positive=False)
        assert_true(coefs.min() < 0)

        alpha, active, coefs = \
            linear_model.lars_path(diabetes['data'], diabetes['target'],
                                   return_path=True, method=method,
                                   positive=True)
        assert_true(coefs.min() >= 0)

# now we gonna test the positive option for all estimator classes

default_parameter = {'fit_intercept': False}

estimator_parameter_map = {'Lars': {'n_nonzero_coefs': 5},
                           'LassoLars': {'alpha': 0.1},
                           'LarsCV': {},
                           'LassoLarsCV': {},
                           'LassoLarsIC': {}}

def test_estimatorclasses_positive_constraint():
    # testing the transmissibility for the positive option of all estimator
    # classes in this same function here

    for estname in estimator_parameter_map:
        params = default_parameter.copy()
        params.update(estimator_parameter_map[estname])
        estimator = getattr(linear_model, estname)(positive=False, **params)
        estimator.fit(diabetes['data'], diabetes['target'])
        assert_true(estimator.coef_.min() < 0)
        estimator = getattr(linear_model, estname)(positive=True, **params)
        estimator.fit(diabetes['data'], diabetes['target'])
        assert_true(min(estimator.coef_) >= 0)

def test_lasso_lars_vs_lasso_cd_positive(verbose=False):
    # Test that LassoLars and Lasso using coordinate descent give the
    # same results when using the positive option

    # This test is basically a copy of the above with additional positive
    # option. However for the middle part, the comparison of coefficient values
    # for a range of alphas, we had to make an adaptations. See below.

    # not normalized data
    X = 3 * diabetes.data

    alphas, _, lasso_path = linear_model.lars_path(X, y, method='lasso',
                                                   positive=True)
    lasso_cd = linear_model.Lasso(fit_intercept=False, tol=1e-8, positive=True)
    for c, a in zip(lasso_path.T, alphas):
        if a == 0:
            continue
        lasso_cd.alpha = a
        lasso_cd.fit(X, y)
        error = linalg.norm(c - lasso_cd.coef_)
        assert_less(error, 0.01)

    # The range of alphas chosen for coefficient comparison here is restricted
    # as compared with the above test without the positive option. This is due
    # to the circumstance that the Lars-Lasso algorithm does not converge to
    # the least-squares-solution for small alphas, see 'Least Angle Regression'
    # by Efron et al 2004. The coefficients are typically in congruence up to
    # the smallest alpha reached by the Lars-Lasso algorithm and start to
    # diverge thereafter.  See
    # https://gist.github.com/michigraber/7e7d7c75eca694c7a6ff

    for alpha in np.linspace(6e-1, 1 - 1e-2, 20):
        clf1 = linear_model.LassoLars(fit_intercept=False, alpha=alpha,
                                      normalize=False, positive=True).fit(X, y)
        clf2 = linear_model.Lasso(fit_intercept=False, alpha=alpha, tol=1e-8,
                                  normalize=False, positive=True).fit(X, y)
        err = linalg.norm(clf1.coef_ - clf2.coef_)
        assert_less(err, 1e-3)

    # normalized data
    X = diabetes.data
    alphas, _, lasso_path = linear_model.lars_path(X, y, method='lasso',
                                                   positive=True)
    lasso_cd = linear_model.Lasso(fit_intercept=False, normalize=True,
                                  tol=1e-8, positive=True)
    for c, a in zip(lasso_path.T[:-1], alphas[:-1]):  # don't include alpha=0
        lasso_cd.alpha = a
        lasso_cd.fit(X, y)
        error = linalg.norm(c - lasso_cd.coef_)
        assert_less(error, 0.01)