#---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# Capacitance calculation 
#
# Provided the path where the files for integrated_density for electrons and holes are stored, it computes the total integrated charge only for the regions
# where are present in both of them. For each of the common regions, linear- and cubic-splines interpolation of the total integrated charge as function of the bias
# is performed. The capacitances are computed as the first derivative of the corresponding interpolated total integrated charge for both cases (linear and cubic). 
# The output file is called "capacitance.dat" and is stored in a subdirectory (/PostProcessing/Capacitance ) within the output Folder of the simulation.
# If the flag -p is present in the command line, the interpolated and capacitance curves as function of bias will be plotted using Matplotlib).
# As default, the bias used for the C-V plot is the first common bias to both integrated density charge files. It is possible to set a another reference bias by
# specifying a substring of the contact name in the input file (using -b1 in the command line) or the difference between two contacts(using -b1 and -b2 in the command 
# line
#
# Script for generation of file with C-V curve from output files from nextnano++
#
# Input: 
# * File "integrated_density_electron.dat"                 # output file from nextnano++
# * File "integrated_density_hole.dat"                     # output file from nextnano++
# * output                                                 # path of the output folder where the files above are stored     
# * bias1                                                  # Substring of the contact that will be used as reference. When not specified the first common
#                                                          # contact of both integrated_density files will be used.
# * bias2                                                  # Substring of a second contact that will be used as reference. The final C-V will be calculated as function of
#                                                          # the voltage given by bias1 - bias 2. If bias1 was not specified, bias2 will be ignored.
# * the flag -p                                            # if present in the command line, the total integrated charge and interpolated C-V curves will be plotted using Matplotlib
#
# Developed by Maria Cecilia da Silva Figueira by nextnano GmbH - version 1.0 (November 4th, 2021)
#---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
                                                 
import argparse, textwrap
import numpy as np
import matplotlib.pyplot as plotter

from os import listdir, makedirs
from os.path import join, exists, dirname
from scipy import interpolate

OUTPUT_DIR = ""
ELEMENTARY_CHARGE = 1.602176487e-19 # SI units (C)   

#---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# Functions

#Identification of common labels in the header of the integrated density files.
#Depending of the number of reference biases provided, verifies if they are present in both files.
#In the case all reference bias and at least one common region, where the charge was integrated, is present
#in both files, then a boolean variable validates the header.
#In case the referenceBiasList is empty, the first bias present in both files is taken as reference.
#As output, the label of the capacitance file and a list of pairs of indexes. Each pair is related to a different region where the 
#charge was integrated. It expresses which column of the integrated electron density and in the integrated hole density files is
#related to the same region. Only regions that the charge was integrated for both kind of carriers are taken into account.

def analyzeHeader(referenceContact,headerElectronList, headerHoleList):
    correspondenceReferenceBiasColumns = []
    correspondenceAllCommonBiasColumns = [] 
    correspondenceChargeColumns = []    
    validReferenceContact = True                              # By default referenceContact is supposed to have the right information
    newLabel = ''
    
    [firstContact,secondContact] = ['','']
 
    #Validation of the number of contacts in the referenceContact list ( only 0, 1 or 2 are allowed) 
    numReferenceContacts = len(referenceContact)

    if numReferenceContacts > 2:
        validReferenceContact = False
    else:
        if numReferenceContacts > 0:
            firstContact = referenceContact[0]
        if numReferenceContacts > 1:
            secondContact = referenceContact[1]

    #Find in the header of the integrated density files the labels related with bias
    commonBiasList =  [ label for label in headerElectronList if label in headerHoleList and "_bias" in label]

    
    #Verify if there is some common contact in both integrated density files and if the references are present in this list
    numCommonBias = len(commonBiasList)
    if numCommonBias == 0:
        validReferenceContact = False
        print("WARNING: There is not any common contact to the integrated density files for electrons or holes. Please check your input file.")

    if numCommonBias < numReferenceContacts:
        validReferenceContact = False
        print("WARNING: At least one of the reference contacts is not common to the integrated density files for electrons or holes. Please check your input file.")
        
    #Build the header of the capacitance file, and maps the position of the common labels in the header of the integrated density files for electron and holes. A correspondence of the 
    #reference bias is also generated. (These maps are essencial for combining the results of two different files)
    if validReferenceContact:  
        # when no reference bias is specified the first common bias in the list will be taken as reference  
        if numReferenceContacts == 0:
            newLabel = commonBiasList[0]
            correspondenceReferenceBiasColumns.append([headerElectronList.index(newLabel),headerHoleList.index(newLabel)] )

        # when one or two biases are specified
        if numReferenceContacts in [1,2]:
            firstContactLabels = [ label for label in commonBiasList if referenceContact[0] in label ]
            if len(firstContactLabels)==0:
                validReferenceContact = False
                print("WARNING: There is not a common contact containing the substring \"" + referenceContact[0] + "\" in the integrated density files for electrons or holes. Please check your input file.")
            else:
                if len(firstContactLabels)>1:
                    validReferenceContact = False
                    print("WARNING: There are more than one contact containing the substring \"" + referenceContact[0] + "\" in, at least, one of the integrated density files for electrons or holes. Please check your input file or choose an unique contact name as the first reference contact.")
                else:
                    newLabel= firstContactLabels[0] 
                    correspondenceReferenceBiasColumns.append([headerElectronList.index(newLabel),headerHoleList.index(newLabel)] )   
 
        # when exactly two reference biases are specified
        if numReferenceContacts == 2 and validReferenceContact:
            secondContactLabels = [ label for label in commonBiasList if referenceContact[1] in label ]       
            if len(secondContactLabels)==0:
                validReferenceContact = False
                print("WARNING: There is not contact called " + referenceContact[1] + " in, at least, one of the integrated density files for electrons or holes.Please check your input file")    
            else:
                if len(secondContactLabels)>1:
                    validReferenceContact = False
                    print("WARNING: There is more than one contact containing the substring \"" + referenceContact[1] + "\" in, at least, one of the integrated density files for electrons or holes. Please check your input file or choose an unique contact name as the second reference contact.")
                else:
                    newLabel= 'V_' + newLabel.split("_bias")[0] + ',' + secondContactLabels[0].replace('_bias','')
                    correspondenceReferenceBiasColumns.append([headerElectronList.index(secondContactLabels[0]),headerHoleList.index(secondContactLabels[0])] )                             

    #generate the other maps only if was successful in the previous part    
    if validReferenceContact:
        commonDensityList =  [ label for label in headerElectronList if label in headerHoleList and "_bias" not in label]
        numCommonDensityRegions = len(commonDensityList)                              
        if numCommonDensityRegions < 1:
            validReferenceContact = False
            print("WARNING: There is not a common region of integration for electrons and holes. Please, check your input file.")
        else:
            for label in commonDensityList:
                newLabel = newLabel + " Capacitance_" + label
                correspondenceChargeColumns =  [ [headerElectronList.index(label),headerHoleList.index(label)] for label in commonDensityList]     
 
        correspondenceAllCommonBiasColumns =  [ [headerElectronList.index(label),headerHoleList.index(label)] for label in commonBiasList]  
 
    return [validReferenceContact, newLabel, correspondenceReferenceBiasColumns, correspondenceAllCommonBiasColumns, correspondenceChargeColumns]
   

#Compute capacitances of the integration regions that are common to the integrated density files
def computeAllCapacitances(header, data, outputFolder, plotResults):
    outputDir = join(outputFolder, 'PostProcessing/Capacitance')
    if not exists(outputDir):
        makedirs(outputDir)

    outputFile = open(join(outputDir,"capacitance.dat"),"w")

    #directory where the saved images will be stored
    plotDir = dirname(outputFile.name).replace('/','\\') + '\\'
    print("\nThe figures generated in Matplotlib will be stored in the directory \"%s\"" % plotDir )

    #Number of common regions of integration
    numCapacitances = len(data[0][1])
    
    #Build a vector with voltage to be plotted in the x-axis
    bias = [row[0] for row in data]  
    numPoints = len(bias)
  
    biasLabel = header[0]
    outputFile.write(biasLabel)
    
    #Build the header of the final capacitance file writing two columns for each integration region, according to a linear or cubic splines interpolation
    for item in header[1:]:
        outputFile.write(" " + item.replace('[',"_linear[") + " " + item.replace('[',"_cubic["))
    outputFile.write('\n')       
    
    # values of voltage to be used in the final C-V to be stored
    xnew = np.linspace(bias[0],bias[numPoints-1],num=numPoints, endpoint=True)    

    allCapacitanceResults = []
    for index in range(0,numCapacitances):
        capacitanceLabel = header[index+1]      
        
        #Build a vector with all values of the total integrated charge in data and interpolates using linear or cubic splines
        totalIntegratedCharge = [row[1][index] for row in data]
        interpolatedCharge = interpolate.splrep(bias,totalIntegratedCharge, k=1)
        interpolatedChargeCubic = interpolate.splrep(bias,totalIntegratedCharge, k=3)   #interp1d(bias,totalIntegratedCharge, kind='cubic')
        
        totalIntegratedChargeLinear = interpolate.splev(xnew,interpolatedCharge,der=0)
        totalIntegratedChargeCubic =  interpolate.splev(xnew,interpolatedChargeCubic,der=0)      

        # computing capacitance from the point of the interpolated charge
        capacitanceLinear = interpolate.splev(xnew,interpolatedCharge,der=1)
        capacitanceCubic =  interpolate.splev(xnew,interpolatedChargeCubic,der=1)        

        if plotResults:
            # total integrated charge from interpolated function
            totalIntegratedChargeLinear = interpolate.splev(xnew,interpolatedCharge,der=0)
            totalIntegratedChargeCubic =  interpolate.splev(xnew,interpolatedChargeCubic,der=0)         
        
            #Format the label of the x-axis
            if ',' in biasLabel:
                biasAxis  = '$' + biasLabel.replace('[','} (').replace(']',')').replace('_','_{',1) + '$'
            else:
                biasAxis = '$V_{' + biasLabel.replace('_bias[','} (').replace(']',')').replace('_','_{',1) + '$'

            #Plot total integrated charge density
            plotter.figure(1 + 2*index) 

            plotter.xlabel(biasAxis)
            plotter.ylabel('Total integrated charge ($C/cm^2$)')
            plotter.plot(bias, totalIntegratedCharge,'bo',xnew,totalIntegratedChargeLinear,'r-',xnew,totalIntegratedChargeCubic,'g--')
            plotter.legend(['data','linear','cubic'],loc='best') 
            figureName = capacitanceLabel.split('[')[0].replace("Capacitance","total_integrated_charge_density") + ".png"
            plotter.savefig(join(plotDir, figureName))            
        
            #Plot capacitance
            plotter.figure(2 + 2*index) 
            plotter.xlabel(biasAxis)
            plotter.ylabel('Capacitance ($\mu F/cm^2$)')

            plotter.plot( xnew,capacitanceLinear,'r-',xnew,capacitanceCubic,'g--')
            if numCapacitances == 1:
                plotter.legend(['linear', 'cubic','linear2','cubic2'],loc='best')       
            else:
                plotter.legend([capacitanceLabel + '_linear', capacitanceLabel + '_cubic'],loc='best')
                
            figureName = capacitanceLabel.split('[')[0]+ ".png"
            plotter.savefig(join(plotDir, figureName))
            plotter.show()
            
        allCapacitanceResults.append([capacitanceLinear,capacitanceCubic])
            
    #Print out the C-V curve in the subfolder PostProcessing/Capacitance
    for bias in range(0,len(xnew)):
        result = str(xnew[bias])
        for region in allCapacitanceResults:
            result = result + " " + str(region[0][bias]) + " " + str(region[1][bias])
        outputFile.write(result + '\n')
        
    outputFile.close()
    
    return
 
 
def compactData(index,biasColumns,densityColumns,line):
    biasColumnList = [item[index] for item in biasColumns]
    densityColumnList = [item[index] for item in densityColumns]
    
    biasList= [float(line[column]) for column in biasColumnList]
    densityList= [float(line[column]) for column in densityColumnList]
    
    return [biasList, densityList]   

      
def generateCapacitanceFile(referenceBiasList, outputFolder, plotResults):
    electronIntegratedDensityFilename = join(outputFolder, "integrated_density_electron.dat")
    holeIntegratedDensityFilename     = join(outputFolder, "integrated_density_hole.dat")
    
    generateOutputFile = True
    
    if not exists(electronIntegratedDensityFilename):
        print("WARNING: The files with integrated electron density is missing in the directory " + outputFolder)
        generateOutputFile = False
        
    if not exists(holeIntegratedDensityFilename):
        print("WARNING: The files with integrated hole density is missing in the directory " + outputFolder)
        generateOutputFile = False        

    if  generateOutputFile:
        
        electronIntegratedDensityFile = open(electronIntegratedDensityFilename, "r")
        holeIntegratedDensityFile     = open(holeIntegratedDensityFilename, "r")
        
        headerElectron = electronIntegratedDensityFile.readline().strip().split(' ')
        headerHole = holeIntegratedDensityFile.readline().strip().split(' ')
        
        [validHeader, newHeader, correspondenceReferenceBiasColumns, correspondenceBiasColumns, correspondenceCapacitanceColumns] = analyzeHeader(referenceBiasList,headerElectron, headerHole)
              
        if validHeader:
              
            electronDensityList = []
            holeDensityList = []            
          
            electronWholeBiasList =[]
            electronWholeDensityList =[]
            
            validData = []
            
            #Read the whole integratedElectronDensity file and generates a vector with all biases and regions of integrations
            for line in electronIntegratedDensityFile:
                lineElectronList = line.strip().split(' ')       
                
                [electronBiasList, electronDensityList] = compactData(0,correspondenceBiasColumns, correspondenceCapacitanceColumns, lineElectronList)

                electronWholeBiasList.append(electronBiasList)
                electronWholeDensityList.append(electronDensityList)

            #Read the each line of integratedHoleDensity file         
            for line in holeIntegratedDensityFile:
                lineHoleList = line.strip().split(' ') 
                
                #Generate a vector with all biases and regions of integrations
                [holeBiasList, holeDensityList] = compactData(1,correspondenceBiasColumns, correspondenceCapacitanceColumns, lineHoleList)


                #Filter only the lines that all biases are the same in both files
                if holeBiasList in electronWholeBiasList:
                    lineInElectronList = electronWholeBiasList.index(holeBiasList)
                    
                    totalIntegratedChargeList = []
                    
                    #Compute the totalIntegrated charge for each common region
                    for electronDensity, holeDensity in zip(electronWholeDensityList[lineInElectronList],holeDensityList):
                        totalIntegratedChargeList.append((electronDensity - holeDensity)* ELEMENTARY_CHARGE*1e6)                      
                    
                    #Stores only the referenceBias and integratedCharge of all regions in validData
                    if len(correspondenceReferenceBiasColumns)==1:
                        referenceBias = holeBiasList[correspondenceReferenceBiasColumns[0][1]]
                    else:
                        if len(correspondenceReferenceBiasColumns) == 2:
                            referenceBias = holeBiasList[correspondenceReferenceBiasColumns[0][1]] - holeBiasList[correspondenceReferenceBiasColumns[1][1]]
                       
                    validData.append([referenceBias,totalIntegratedChargeList])                            
                    
            headers = " ".join(newHeader.split(" "))
            print("\nThe information to be included in the header of the capacitance file is concerning to : " + headers)
            
            computeAllCapacitances(newHeader.split(),validData,outputFolder, plotResults)
               
        electronIntegratedDensityFile.close()
        holeIntegratedDensityFile.close()
    return


#---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# Parser

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, description = textwrap.dedent('''\

    Capacitance calculation - nextnano GmbH 
    
    Provided the path where the files of integrated density for electrons and holes are stored, it computes the total integrated charge only for the regions where are present in both of them. For each of the common regions, linear- and cubic-splines interpolation of the total integrated charge as function of the bias
    is performed. The capacitances are computed as the first derivative of the corresponding interpolated total integrated charge for both cases (linear and cubic).  The output file is called "capacitance.dat" and is stored in a subdirectory (/PostProcessing/Capacitance ) within the output Folder of the simulation.
    If the flag -p is present in the command line, the interpolated and capacitance curves as function of bias will be plotted using Matplotlib). As default, the bias used for the C-V plot is the first common bias to both integrated density charge files. It is possible to set a another reference bias by
    specifying a substring of the contact name in the input file (using -b1 in the command line) or the difference between two contacts(using -b1 and -b2 in the command line.
    
    Script for generation of file with C-V curve from output files from nextnano++
    
    Minimum Requirements: 
    * File "integrated_density_electron.dat"                 # output file from nextnano++
    * File "integrated_density_hole.dat"                     # output file from nextnano++
    * output                                                 # path of the output folder where the files above are stored  
    
'''))

parser = argparse.ArgumentParser(add_help=False)

parser.add_argument('-h', '--help', action='help', default=argparse.SUPPRESS, help=textwrap.dedent('''\
   Show this message and exit.
'''))

parser.add_argument('-o','--output', type = str, required= True,  help=textwrap.dedent('''\
   Directory where the simulation results are stored.
'''))
parser.add_argument('-b1','--bias1', type = str, required= False, help=textwrap.dedent('''
   Substring of the contact that will be used as reference. When not specified the first common
   contact in the integrated_density files (for electrons and holes) will be used. Example: for
   GateContact_bias as reference use: -b1 Gate
'''))

parser.add_argument('-b2','--bias2', type = str, required= False, help=textwrap.dedent('''\
   Substring of a second contact that will be used as reference. The final C-V will be calculated as function of
   the voltage given by bias1 - bias 2. If bias1 was not specified, bias2 will be ignored. Example: for
   reference bias as the difference of GateContact_bias - SubstrateContact_bias, use: -b1 Gate -b2 Substrate
'''))
                                                                                            
parser.add_argument('-p','--plot', default=False, action="store_true", help= ('''\
   Plot the total integrated charge and capacitances using Matplotlib
'''))     

arguments = parser.parse_args()

if arguments.output:
    outputFolder = arguments.output

referenceBias = []    
if arguments.bias1:
    referenceBiasText = "substrings \"" + arguments.bias1 + "\""
    if arguments.bias2:
        referenceBiasText = "substrings \"" + arguments.bias1 + "\" and \"" + arguments.bias2+ "\""
        referenceBias = [arguments.bias1, arguments.bias2]
    else:
        referenceBias = [arguments.bias1]
else:
    if arguments.bias2:
        referenceBiasText = "substring \"" + arguments.bias2 + "\" . It will be ignored, because bias1 was not specified in the command line"
    else:
        referenceBiasText = "no reference bias specified in the command line"    

plotResults = True

if arguments.plot:
    plotResults = True
else:
    plotResults = False

print("\nArguments of the command line\n")
print("Output folder   " + outputFolder)
print("Reference Bias  " + referenceBiasText)
print("Plot option     " + str(plotResults) + "\n\n")


#---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

#main

generateCapacitanceFile(referenceBias,outputFolder, plotResults)