#!/usr/bin/env python3
#
#  Purpose: Combine the AK recreated equation with the CO recreated equation
#  History:
#           Jan. 2022      Kochenash    - newly created for LMP
#  Usage:
#         example: eqn_combine.py ${EQNFILE}_co_rec ${EQNFILE}_ak_rec --equation_out ${EQNFILE}_com
#            
import argparse
from math import ceil

class equations():
    def __init__(self):
        # initalize data structores for holding equation file strings
        # Every entry in tands has an associated list of entries in stations and tors
        self.tand_strings = []
        self.station_strings = []
        self.tor_strings = []
        self.heads = []
        
        
    def read(self,file):
        
        def read_tand(lines,start_at=0):
            #prepare to read by setting up some starting settings
            tands = True
            station = False
            tors = False
            ntands = int(lines[start_at+1])
            endtands = ntands
            
            tand_ix = len(self.tand_strings)
            self.tor_strings.append([])
            self.station_strings.append([])
            eqn_ix = len(self.tor_strings[tand_ix])
        
            self.heads.append(lines[0])
        
            for ix, l in enumerate(lines[start_at+1:]):

                if tands:
                    if ix == endtands-ntands:
                        temp_list = []
                        self.station_strings[tand_ix].append([])
                        self.tor_strings[tand_ix].append([])
                    if ix < endtands:
                        temp_list.append(l)
                    else:
                        temp_list.append(l)
                        temp_tup = tuple(temp_list)
                        self.tand_strings.append(temp_tup)
                        tands = False
                        stations = True

                elif stations:
                    self.station_strings[tand_ix][eqn_ix].append(l)
                    if l.strip().endswith("99999999"):                        
                        stations = False
                        tors = True
                        ntors=int(lines[start_at+ix+2]) # the next line
                        m = ceil(ntands/10.0)
                        endtors = ix+ntors+(ntors*m)+(3*m)+1 #Three extra 'coef' lines and one line with number of preds when more than 10 predictors takes more lines so multiply by m
                        
                
                elif tors:
                    self.tor_strings[tand_ix][eqn_ix].append(l)
                    if ix == endtors:
                        tors = False
                        if lines[start_at+ix+2].strip() == "99999999":                  
                            break
                        else:
                            stations = True
                            eqn_ix+=1
                            self.station_strings[tand_ix].append([])
                            self.tor_strings[tand_ix].append([])
            return start_at+ix+2
            
        with open(file) as f:
            lines = f.readlines()
        head = lines[0]
        last = read_tand(lines)
        
        
        while last+1 < len(lines):
            last = read_tand(lines,start_at=last+1)
            
            
    def merge(self):
        if len(set(self.heads)) > 1:
            raise ValueError("equation headers do not match")
        if not len(self.tand_strings) == len(self.station_strings) == len(self.tor_strings):
            raise ValueError("predictand, station, and predictor count error")
        self.eqns = {}
        for tand, stns, tors in zip(self.tand_strings,self.station_strings,self.tor_strings):
            if tand not in self.eqns.keys():
                self.eqns[tand] = [stns,tors]
            else:
                for s,t in zip(stns,tors):
                    self.eqns[tand][0].append(s)
                    self.eqns[tand][1].append(t)
    
    def write_merged(self,out_file):
        with open(out_file,'w') as f:
            for k,v in self.eqns.items():
                f.write(self.heads[0])
                
                for l in k:
                    f.write(l)
                    
                for s,t in zip(v[0],v[1]):
                    for l in s:
                        f.write(l)
                        
                    for l in t:
                        f.write(l)

                f.write(' 99999999\n')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Merge multiple MOS2k equation files into one")
    parser.add_argument("equations", nargs = '+', type=str)
    parser.add_argument("--equation_out", default='out_eqn', type=str)
    args = parser.parse_args()

    eqns = equations()
    for equation in args.equations:
        eqns.read(equation)
    eqns.merge()
    eqns.write_merged(args.equation_out)