#!/usr/bin/python

import sys,getopt,commands

from model.symbolic import *
from model.model import *
from solver.solver import *
from misc.calculus import *
from compilator.compilator import *
import re

def dynare_import(filename,ignore_shocks=False):
    otxt = file(filename).read()

    regex = re.compile("(.*)//(.*)")
    def removecomment(line):
        res = regex.search(line)
        if res:
            l = res.groups(1)[0]
            return(l)
        else:
            return line
    txt = str.join("\n",map(removecomment,otxt.split("\n")))
    
    txt = txt.replace("\n","")
    
    
    var_lines = re.compile("var([^;]*);").search(txt).group(1)
    variables_names = var_lines.replace(" ","").split(",")
    set_variables(variables_names)

    
    shocks_lines = re.compile("varexo([^;]*);").search(txt).group(1)
    shocks_names = shocks_lines.replace(" ","").split(",")
    set_shocks(shocks_names)
    
    parameters_lines = re.compile("parameters([^;]*);").search(txt).group(1)
    parameters_names = parameters_lines.replace(" ","").split(",")
    set_parameters(parameters_names)
    print(parameters)
    
    model_lines = re.compile("model;(.*?);end;").search(txt).group(1)
    equations_lines = model_lines.split(';')
    equations = []
    for teq in equations_lines:
        print("teq : " + str(teq))
        teq = teq.replace('^',"**")
        teqlhs,teqrhs = teq.split("=")
        eqlhs = eval(teqlhs)
        eqrhs = eval(teqrhs)
        equations.append(Equation(eqlhs,eqrhs))
        #eq = eval(teq,)
    
    param_def_lines = re.compile("(?:parameters)(?:[^;]*);(.*);model;").search(txt).group(1)
    param_def_lines = param_def_lines.replace(" ","").replace("=",":").replace("^","**")
    param_def = "{ " + str.join(",",param_def_lines.split(";")) + "}"
    parameters_values = eval(param_def)
    
    val_def_lines = re.compile("initval;(.*?);end;").search(txt).group(1)
    val_def_lines = val_def_lines.replace(" ","").replace("=",":").replace("^","**")
    val_def = "{ " + str.join(",",val_def_lines.split(";")) + "}"
    init_values = eval(val_def)
    
    covariances = zeros([len(shocks),len(shocks)])
    if not ignore_shocks:
        shocks_def_lines = re.compile("shocks;(.*?);end;").search(txt).group(1)
        shocks_def_lines = shocks_def_lines.split(";")
        # doesn't work when specifying correlations
        var_order = {} # hash mapping variables to their index
        for i in range(len(shocks)):
            var_order[shocks[i]] = i
        
        regex = re.compile("var (.*?)(|,.*)=(.*)")
        for l in shocks_def_lines:
            lhs,mhs,rhs = regex.search(l).groups()
            lv = eval(lhs)
            if mhs == "":
                mv = lv
            else:
                mv = eval(mhs)
            covariances[var_order[lv],var_order[mv]] = float(rhs)

    
    frame = inspect.currentframe().f_back
    frame.f_globals['variables'] = variables
    frame.f_globals['exovariables'] = []
    frame.f_globals['parameters'] = parameters
    frame.f_globals['shocks'] = shocks
    frame.f_globals['equations'] = equations
    frame.f_globals['parameters_values'] = parameters_values
    frame.f_globals['init_values'] = init_values
    frame.f_globals['covariances'] = covariances
    
    # add variables to workspace
    for v in (variables+shocks+parameters):
        frame.f_globals[v.name] = v
    
    del frame
    

if __name__ == "__main__":

    def main(argv):
        if len(argv)<1:
            print("not enough argument")
            sys.exit(2)
        
        filetype = ""
        
        filename = argv[-1] # last argument is the filename
        regex_mod = re.compile("(.*)\.mod")
        regex_mod_match = re.match(regex_mod,filename)
        if regex_mod_match:
            filetype = "mod"
            filename_trunc = regex_mod_match.groups()[0]
    
        regex_xml = re.compile("(.*)\.xml")
        regex_xml_match = re.match(regex_xml,filename)
        if regex_xml_match:
            filetype = "xml"
            filename_trunc = regex_xml_match.groups()[0]
        
        if filetype == "":
            print("Unknown filetype")
            sys.exit(2)
        
        # Read options
        options = {"ramsey":False,"portfolio":False}
        short_arg_dict = "hr"
        long_arg_dict = ["help","ramsey","portfolio"]
        try:
            opts, args = getopt.getopt(argv, short_arg_dict, long_arg_dict)
        except getopt.GetoptError:
            usage()
            sys.exit(2)
        for opt, arg in opts:
            if opt in ("-h","--help"):
                usage()
                sys.exit()
            if opt in ("-r","--ramsey"):
                options["ramsey"] = True
            if opt in ("--portfolio"):
                options["portfolio"] = True
                
        # Start the actual work        
        if filetype == "mod":
            process_mod_file(filename_trunc)
        elif filetype == "xml":
            model = read_xml_file(filename_trunc)
            if options["ramsey"]:
                write_ramsey_policy(filename_trunc, options, model)
            if options["portfolio"]:
                write_portfolio_version(filename_trunc, options, model)
        
    def process_mod_file(filename_trunc,options={}):
        '''read mod file ; write xml file'''
        file_path = filename_trunc
        print("Processing modfile : " + file_path + '.mod')
        output_file_path = filename_trunc + '_dd'
        command = "ruby DareDare.rb convert " + file_path + " " + output_file_path
        output = commands.getstatusoutput(command)
        if not output[0]:
            print("Conversion successfull. Written to : " + output_file_path + ".xml")
            sys.exit(2)
    
    def read_xml_file(filename_trunc,options={}):
        '''process xml file ; returns parsed model with symbolic equations'''
        print("Processing xmlfile : " + filename_trunc + ".xml")
        f = file(filename_trunc + ".xml")
        txt = f.read()
        f.close()
        model = Model(filename_trunc)
        model.import_from_xml(ET.XML(txt))
        model.sympify_model()   
        return(model)
    
    def write_ramsey_policy(filename_trunc,options,model):
        #model.export_to_modfile()
        ramsey_model = model.process_ramsey_policy_model()
        f = open(filename_trunc + "_ramsey.mod","w")
        f.write(ramsey_model.export_to_modfile())
        f.close()
        print("Ramsey model written to : " + filename_trunc + "_ramsey.mod")
    
    def write_portfolio_version(filename_trunc,options,model):
        print("Processing portfolios")
        portfolio_model = model.process_portfolio_model()
        f = open(filename_trunc + "_pf.mod","w")
        f.write(portfolio_model.export_to_modfile())
        f.close()
        print("Portfolio model written to : " + filename_trunc + "_pf.mod")
    
    def usage():
        print("Bad arguments")
    
    print(sys.argv)
    main(sys.argv[1:])
