#!/usr/bin/env python3

# (c) Copyright 2020 CORSIKA Project, corsika-project@lists.kit.edu
#
# See file AUTHORS for a list of contributors.
#
# This software is distributed under the terms of the GNU General Public
# Licence version 3 (GPL Version 3). See file LICENSE for a full version of
# the license.


import pickle, sys, itertools



# loads the pickled particle_db (which is an OrderedDict)
def load_particledb(filename):
    with open(filename, "rb") as f:
        particle_db = pickle.load(f)
    return particle_db



# 
def read_qgsjetII_codes(filename, particle_db):
    with open(filename) as f:
        for line in f:
            line = line.strip()
            if len(line)==0 or line[0] == '#':
                continue
            line = line.split('#')[0]
            print (line)
            identifier, model_code, xsType = line.split()
            try:
                particle_db[identifier]["qgsjetII_code"] = int(model_code)
                particle_db[identifier]["qgsjetII_xsType"] = int(xsType)
            except KeyError as e:
                raise Exception("Identifier '{:s}' not found in particle_db".format(identifier))


            

# generates the enum to access qgsjetII particles by readable names
def generate_qgsjetII_enum(particle_db):
    output = "enum class QgsjetIICode : int8_t {\n"
    for identifier, pData in particle_db.items():
        if pData.get('qgsjetII_code') != None:
            output += "  {:s} = {:d},\n".format(identifier, pData['qgsjetII_code'])
    output += "};\n"
    return output



# generates the look-up table to convert corsika codes to qgsjetII codes
def generate_corsika2qgsjetII(particle_db):    
    string = "std::array<QgsjetIICodeIntType, {:d}> constexpr corsika2qgsjetII = {{\n".format(len(particle_db))
    for identifier, pData in particle_db.items():
        modelCode = pData.get("qgsjetII_code", 0)
        string += "  {:d}, // {:s}\n".format(modelCode, identifier if modelCode else identifier + " (not implemented in QGSJETII)")
    string += "};\n"
    return string
    


# generates the look-up table to convert corsika codes to qgsjetII codes
def generate_corsika2qgsjetII_xsType(particle_db):    
    string = "std::array<int, {:d}> constexpr corsika2qgsjetIIXStype = {{\n".format(len(particle_db))
    for identifier, pData in particle_db.items():
        modelCodeXS = pData.get("qgsjetII_xsType", -1)
        string += "  {:d}, // {:s}\n".format(modelCodeXS, identifier if modelCodeXS else identifier + " (not implemented in QGSJETII)")
    string += "};\n"
    return string


# generates the look-up table to convert qgsjetII codes to corsika codes    
def generate_qgsjetII2corsika(particle_db) :
    minID = 0
    for identifier, pData in particle_db.items() :
        if 'qgsjetII_code' in pData:
            minID = min(minID, pData['qgsjetII_code'])

    string = "QgsjetIICodeIntType constexpr minQgsjetII = {:d};\n\n".format(minID)

    pDict = {}
    for identifier, pData in particle_db.items() :
        if 'qgsjetII_code' in pData:
            model_code = pData['qgsjetII_code'] - minID
            pDict[model_code] = identifier
    
    nPart = max(pDict.keys()) - min(pDict.keys()) + 1
    string += "std::array<corsika::particles::Code, {:d}> constexpr qgsjetII2corsika = {{\n".format(nPart)
    
    for iPart in range(nPart) :
        identifier = pDict.get(iPart, "Unknown")
        qgsID = iPart + minID
        string += "  corsika::particles::Code::{:s}, // {:d} \n".format(identifier, qgsID)
    
    string += "};\n"
    return string

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("usage: {:s} <particle_db.pkl> <qgsjetII_codes.dat>".format(sys.argv[0]), file=sys.stderr)
        sys.exit(1)
        
    print("code_generator.py for QGSJETII")
    
    particle_db = load_particledb(sys.argv[1])
    read_qgsjetII_codes(sys.argv[2], particle_db)
    
    with open("Generated.inc", "w") as f:
        print("// this file is automatically generated\n// edit at your own risk!\n", file=f)
        print(generate_qgsjetII_enum(particle_db), file=f)
        print(generate_corsika2qgsjetII(particle_db), file=f)
        print(generate_qgsjetII2corsika(particle_db), file=f)
        print(generate_corsika2qgsjetII_xsType(particle_db), file=f)