#!/usr/bin/python3
#
'''
Convert test result in the RITdb to CSV format

Usage: RITdb2csv.py [-p | --pivot] [-s | --split] -i <input RITdb file> [-o <output CSV file>]
'''

import os, sys, getopt
import sqlite3 as sql
import csv

PartInfoQuery='''
SELECT  n2.value AS 'Part ID', n3.value AS 'Pass/Fail', n4.value AS 'Test Time', n5.value AS 'Cycle Time', n6.value AS 'Site', n1.value AS 'part result ID', n0.entityID AS 'entityID'
FROM ritdb1 n0
 JOIN ritdb1 n1 ON n0.entityID=n1.entityID AND n1.name='PART_RESULT_EVENT_ORDER'
 JOIN ritdb1 n2 ON n0.entityID=n2.entityID AND n2.name='PART_ID'
 JOIN ritdb1 n3 ON n0.entityID=n3.entityID AND n3.name='PF'
 JOIN ritdb1 n4 ON n0.entityID=n4.entityID AND n4.name='EVENT_TEST_TIME'
 JOIN ritdb1 n5 ON n0.entityID=n5.entityID AND n5.name='EVENT_CYCLE_TIME'
 JOIN ritdb1 n6 ON n0.entityID=n6.entityID AND n6.name='SITE_ID'
 {}
WHERE
 n0.value='PART_RESULT_EVENT'{}
ORDER BY n1.value ASC '''

ResultInfoQuery='''
SELECT  n1.value AS 'Result Number',  n2.value AS'Result Name',  n4.value AS 'Units',  n8.value AS 'U Limit',  n9.value AS 'L Limit',  n3.value AS 'RESULT_ID'
FROM ritdb1 n0
 JOIN ritdb1 n1 ON n0.entityID=n1.entityID AND n1.indexID='0' AND n1.name='RESULT_NUMBER'
 JOIN ritdb1 n2 ON n0.entityID=n2.entityID AND n2.indexID='0' AND n2.name='RESULT_NAME'
 JOIN ritdb1 n3 ON n0.entityID=n3.entityID AND n3.indexID='0' AND n3.name='RESULT_ID'
 JOIN ritdb1 n4 ON n0.entityID=n4.entityID AND n4.indexID='0' AND n4.name='RESULT_UNITS'
 JOIN ritdb1 n5 ON n5.indexID='0' AND n5.name='TEST_SPECIFICATION_NAME'
 JOIN ritdb1 n6 ON n6.indexID='0' AND n6.name='ENTITY_TYPE' AND n6.value='RESULT_LIMIT_SET'
 JOIN ritdb1 n7 ON n6.entityID=n7.entityID AND n5.value=n7.value AND n7.indexID='0' AND n7.name='LIMIT_SET_NAME'
 LEFT JOIN ritdb1 n8 ON n0.entityID=n8.indexID AND n6.entityID=n8.entityID AND n8.name='UL'
 LEFT JOIN ritdb1 n9 ON n0.entityID=n9.indexID AND n6.entityID=n9.entityID AND n9.name='LL'
WHERE
 n0.name='ENTITY_TYPE' AND n0.value='RESULT_INFO'
ORDER BY n3.value ASC '''

ResultsQuery='''
SELECT  n0.value * n3.value AS 'Result',  n0.value2,  n1.value,  n2.value
FROM ritdb1 n0
 JOIN ritdb1 n1 ON n0.entityID=n1.entityID AND n1.name='PART_RESULT_EVENT_ORDER'
 JOIN ritdb1 n2 ON n0.indexID=n2.entityID AND n2.name='RESULT_ORDER'
 JOIN ritdb1 n3 ON n2.entityID=n3.entityID AND n3.name='RESULT_SCALE'
 {}
WHERE
 n0.name='R'{} '''
 
WaferEIDQuery='''
SELECT n0.entityID AS 'WaferEID', n1.value AS 'WaferID'
FROM ritdb1 n0
 JOIN ritdb1 n1 ON n0.entityID=n1.entityID AND n1.name='SUBSTRATE_ID'
Where
 n0.value='SUBSTRATE_EVENT' '''

SubstrateEven_EID = '''JOIN ritdb1 n9 ON n0.entityID=n9.entityID AND n9.name='SUBSTRATE_EVENT_EID' '''

def ritdb2csv(dbCursor, csvFileName, modStr):

#    query the top header
    if modStr: query = PartInfoQuery.format(SubstrateEven_EID, modStr)
    else: query = PartInfoQuery.format('', '')
    dbCursor.execute(query)
    topTable = dbCursor.fetchall()
    nDevice = len(topTable)
    topLabel = []
    topDict = {}   # dict{PART_RESULT_EVENT_ORDER : position}
    for i in range(4): topLabel.append(dbCursor.description[i][0])
    for i in range(nDevice): topDict.update({topTable[i][5]: i})
    
#    query the left table
    dbCursor.execute(ResultInfoQuery)
    leftTable = dbCursor.fetchall()
    nTest = len(leftTable)
    leftLabel = []
    for i in range(5): leftLabel.append(dbCursor.description[i][0])

#    query result data, sort by PART_RESULT_EVENT_ORDER then RESULT_ORDER
    if modStr: query = ResultsQuery.format(SubstrateEven_EID, modStr)
    else: query = ResultsQuery.format('', '')
    query = query + 'ORDER BY n2.value ASC'
    dbCursor.execute(query)

#    open csv writer
    with open(csvFileName, "w") as csv_file:
        csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
  
#    write the top table to csv file
        for i in range(4):
            row = ['','','','','',topLabel[i]]
            for n in range(nDevice): row.append(topTable[n][i])
            csv_writer.writerow(row)

        csv_writer.writerow(leftLabel)

#    Write one row at a time
        data = dbCursor.fetchone()     # get the first row of resultQuery data
        for ltRow in leftTable:
            dataLists = ['' for i in range(nDevice)]
            row = []
            for i in range(5): row.append(ltRow[i])
            row.append('')

            while data != None and ltRow[5] == data[3]:     # leftTable.testEntityID = data.testEntityID
                dataLists[topDict[data[2]]] = data[0]       # device ID is 1 index
                data = dbCursor.fetchone()                    
            row.extend(dataLists)
            csv_writer.writerow(row)
             
### end of ritdb2csv ###

def ritdb2tcsv(dbCursor, csvFileName, modStr):

#    query the top header
    dbCursor.execute(ResultInfoQuery)
    topTable = dbCursor.fetchall()
    nTest = len(topTable)
    topLabel = []
    topDict = {}   # dict{PART_RESULT_EVENT_ORDER : position}
    for i in range(5): topLabel.append(dbCursor.description[i][0])
    for i in range(nTest): topDict.update({topTable[i][5]: i})
    
#    query the left table
    if modStr: query = PartInfoQuery.format(SubstrateEven_EID, modStr)
    else: query = PartInfoQuery.format('', '')
    dbCursor.execute(query)
    leftTable = dbCursor.fetchall()
    nDevice = len(leftTable)
    leftLabel = []
    for i in range(4): leftLabel.append(dbCursor.description[i][0])

#    query result data, sort by testEntityID then device ID
    if modStr: query = ResultsQuery.format(SubstrateEven_EID, modStr)
    else: query = ResultsQuery.format('', '')
    query = query + 'ORDER BY n1.value ASC'
    dbCursor.execute(query)

#    open csv writer
    with open(csvFileName, "w") as csv_file:
        csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
  
#    write the top table to csv file
        for i in range(5):
            row = ['','','','',topLabel[i]]
            for n in range(nTest): row.append(topTable[n][i])
            csv_writer.writerow(row)

        csv_writer.writerow(leftLabel)

#    Write one row at a time
        data = dbCursor.fetchone()     # get the first row of resultQuery data
        for ltRow in leftTable:
            dataLists = ['' for i in range(nTest)]
            row = []
            for i in range(0,4): row.append(ltRow[i])
            row.append('')

            while data != None and ltRow[5] == data[2]:     # leftTable.deviceID = data.deviceID
                dataLists[topDict[data[3]]] = data[0]       # testResultID is 1 index
                data = dbCursor.fetchone()                    
            row.extend(dataLists)
            csv_writer.writerow(row)

### end of ritdb2tcsv ###

if __name__ == '__main__':
    
    usage = ("Usage: %s [-p | --pivot] [-s | --split] -i <input RITdb file> [-o <output CSV file>]" % (sys.argv[0]))

    dbFileName = ''
    csvFileName = ''
    pivot = False
    split = False
    
    argv = sys.argv[1:]   
    try:
        opts, args = getopt.getopt(argv,"hpsi:o:",["pivot","split","ifile=","ofile="])

    except getopt.GetoptError:
        print(usage)
        sys.exit(2)
        
    for opt, arg in opts:
        if opt == '-h':
            print(usage)
            sys.exit()
        elif opt in ("-p", "--pivot"): pivot = True
        elif opt in ("-s", "--split"): split = True
        elif opt in ("-i", "--ifile"): dbFileName = arg
        elif opt in ("-o", "--ofile"): csvFileName = arg
            
    if not bool(dbFileName):    # dbFileName is empty   
        print(usage)
        sys.exit(2)
        
    if bool(csvFileName):       # csvFileName is not empty
        baseName, ext = os.path.splitext(csvFileName)
        if ext: csvFileName = baseName + '{}' + ext
        else: csvFileName = baseName + '{}.csv'
    else:                       # csvFileName is empty
        baseName, ext = os.path.splitext(dbFileName)
        csvFileName = (baseName + '{}.csv')
    
    if not os.path.exists(dbFileName):
        print("RITdb file not found:", dbFileName)
        sys.exit(2)
        
    try:
#    Connect to database
        dbConn = sql.connect(dbFileName)
        dbCursor = dbConn.cursor()
        
        if split: 
            dbCursor.execute(WaferEIDQuery)
            waferEID = dbCursor.fetchall()
            if waferEID:
                for eid in waferEID:
                    string = ' AND n9.value=' + str(eid[0])
                    newFileName = csvFileName.format(f"{eid[1]:0>2}")  # pad one 0
                    if pivot:
                        ritdb2tcsv(dbCursor, newFileName, string)
                    else:
                        ritdb2csv(dbCursor, newFileName, string)
                dbCursor.close()
                sys.exit(0)
            else:
                dbCursor.close()
                print("No wafer info found in RITdb:", dbFileName)
                sys.exit(2)
                
        newFileName = csvFileName.format('')
        if pivot:
            ritdb2tcsv(dbCursor, newFileName, '')
        else:
            ritdb2csv(dbCursor, newFileName, '')
            
    except sql.Error as error:
        print("Failed to read data from RITdb;", error)

# Close database connection
    finally:
        dbConn.close()    # Close database connection