# Copyright © 2018-2026 Dynare Team
#
# This file is part of Dynare.
#
# Dynare is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Dynare is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Dynare.  If not, see <https://www.gnu.org/licenses/>.

import re

from pygments.lexer import Lexer, RegexLexer, bygroups, do_insertions, words, include
from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
    Number, Literal, Punctuation, Generic, Whitespace

__all__ = ['DynareLexer']

class DynareLexer(RegexLexer):

    name = 'Dynare'
    aliases = ['dynare']
    filenames = ['*.mod']

    flags = re.MULTILINE | re.IGNORECASE

    # We follow the same categorization as in scripts/dynare.el

    dynare_statements = (
        "var", "varexo", "varexo_det", "trend_var", "log_trend_var",
        "predetermined_variables", "parameters", "model_local_variable", "heterogeneity_dimension",
        "model_info",
        "estimation", "set_time", "data", "varobs", "rplot",
        "osr_params", "osr", "dynatype", "dynasave", "model_comparison", "change_type",
        "load_params_and_steady_state", "save_params_and_steady_state",
        "write_latex_dynamic_model", "write_latex_static_model",
        "write_latex_original_model", "write_latex_steady_state_model", "steady",
        "check", "simul", "stoch_simul", "var_model", "trend_component_model",
        "var_expectation_model", "pac_model", "dsample", "planner_objective",
        "ramsey_model", "ramsey_policy", "evaluate_planner_objective", "occbin_setup",
        "occbin_solver", "occbin_write_regimes", "occbin_graph",
        "discretionary_policy", "identification", "bvar_density", "bvar_forecast",
        "bvar_irf", "sensitivity", "initval_file", "histval_file",
        "forecast", "shock_decomposition", "realtime_shock_decomposition",
        "plot_shock_decomposition", "initial_condition_decomposition",
        "squeeze_shock_decomposition", "sbvar", "ms_estimation", "ms_simulation",
        "ms_compute_mdd", "ms_compute_probabilities", "ms_forecast", "ms_irf",
        "ms_variance_decomposition", "conditional_forecast",
        "plot_conditional_forecast", "method_of_moments", "markov_switching", "svar",
        "svar_global_identification_check", "external_function", "calib_smoother",
        "model_diagnostics", "extended_path", "smoother2histval",
        "perfect_foresight_setup", "perfect_foresight_solver",
        "perfect_foresight_with_expectation_errors_setup",
        "perfect_foresight_with_expectation_errors_solver", "compilation_setup",
        "model_remove", "model_options", "var_remove",
        "heterogeneity_load_steady_state", "heterogeneity_compute_steady_state",
        "heterogeneity_solve", "heterogeneity_simulate",
        "resid", "database", "std", "corr",
        "prior_function", "posterior_function", "end"
    )

    dynare_statements_like = (
        "stderr", "skew", "values", "periods", "scales", "restriction", "exclusion",
        "upper_cholesky", "lower_cholesky", "equation", "name", "bind", "relax", "error_bind",
        "error_relax", "add", "multiply", "target", "auxname_target_nonstationary",
        "component", "growth", "auxname", "kind", "weights", "exogenize", "endogenize"
    )

    dynare_blocks = (
        "model", "steady_state_model", "initval", "endval", "histval",
        "filter_initial_state", "shocks", "heteroskedastic_shocks", "shock_groups",
        "init2shocks", "mshocks", "shock_paths", "estimated_params", "epilogue", "priors",
        "estimated_params_init", "estimated_params_bounds",
        "estimated_params_remove", "osr_params_bounds", "observation_trends",
        "deterministic_trends", "optim_weights", "homotopy_setup",
        "conditional_forecast_paths", "svar_identification", "moment_calibration",
        "irf_calibration", "ramsey_constraints", "generate_irfs", "matched_moments",
        "occbin_constraints", "model_replace", "pac_target_info", "matched_irfs",
        "matched_irfs_weights", "perfect_foresight_controlled_paths", "filter_tunes",
        "verbatim"
    )

    dynare_functions = (
        "expectation", "var_expectation", "pac_expectation", "pac_target_nonstationary",
        "exp", "diff", "log",
        "ln", "log10", "sin", "cos", "tan", "asin", "acos", "atan",
        "sinh", "cosh", "tanh", "asinh", "acosh", "atanh", "sqrt", "cbrt", "abs",
        "sign", "max", "min", "normcdf", "normpdf", "erf", "erfc", "steady_state", "sum"
    )

    dynare_constants = (
        "inf", "nan"
    )

    # The macro keywords are separated in 4 sets depending on their type
    # NB: “for” is both a command and an operator
    dynare_macro_commands = (
        "line", "include", "includepath", "define", "echo", "error", "if", "ifdef", "ifndef",
        "elseif", "else", "endif", "for", "endfor", "echomacrovars"
    )
    dynare_macro_constants = (
        "true", "false", "nan", "inf"
    )
    dynare_macro_functions = (
        "exp", "log", "ln", "log10", "sin", "cos", "tan", "asin", "acos",
        "atan", "sqrt", "cbrt", "sign", "max", "min", "floor", "ceil", "trunc", "mod", "sum",
        "erf", "erfc", "gamma", "lgamma", "round", "length", "normpdf", "normcdf", "isempty",
        "isboolean", "isreal", "isstring", "istuple", "isarray", "defined"
    )
    dynare_macro_word_operators = (
        "for", "in", "when", "bool", "real", "string", "tuple", "array"
    )

    tokens = {
        'root': [
            (r'(%|//).*', Comment.Single),
            (r"/\*", Comment.Multiline, "multiline_comment"),

            (r"@#", Punctuation, "macro_directive"),
            (r"@{", Punctuation, "macro_expansion"),

            (words(dynare_statements + dynare_statements_like + dynare_blocks
                   + dynare_functions, prefix=r'\b', suffix=r'\b'), Keyword.Reserved),

            (words(dynare_constants, prefix=r'\b', suffix=r'\b'), Keyword.Constant),


            # The name/identifier rule must come after all rules that match text keywords
            (r'[a-z_][a-z0-9_]*', Name),

            # The date rule must come before the integer rule, otherwise the year component will be matched as an integer
            (r'-?[0-9]+([ya]|m([1-9]|1[0-2])|q[1-4]|[sh][12])', Literal.Date),
            (r'((([0-9]*\.[0-9]+)|([0-9]+\.))([ed][-+]?[0-9]+)?)|([0-9]+[ed][-+]?[0-9]+)', Number.Float),
            (r'\d+', Number.Integer),

            (r"'[^']*'", String.Single),

            # Must come before the “>” operator
            (r'>>', Punctuation), # For MATLAB prompts

            (r'(-|\+|\*|\/|\^)', Operator),
            (r'(==|<=|>=|!=|<|>)', Operator),
            (r'⟂', Operator),
            (r'(=|&|\||&&|\|\||!|:|~|~=)', Operator), # MATLAB/Octave-specific

            (r'[\[\](){}@.,;#\$]', Punctuation),

            (r'\s+', Whitespace)
        ],

        "multiline_comment": [
            # Comment body, anything except * or /
            (r"[^*/]+", Comment.Multiline),
            # Nested multiline comment: push same state
            (r"/\*", Comment.Multiline, "#push"),
            # End of multiline comment: pop state
            (r"\*/", Comment.Multiline, "#pop"),
            # Single * or / inside comment
            (r"[*/]", Comment.Multiline),
        ],

        "macro_common" : [
            (r'//.*', Comment.Single),

            (words(dynare_macro_functions, prefix=r'\b', suffix=r'\b'), Keyword.Reserved),
            (words(dynare_macro_constants, prefix=r'\b', suffix=r'\b'), Keyword.Constant),
            (words(dynare_macro_word_operators, prefix=r'\b', suffix=r'\b'), Operator.Word),

            # The name/identifier rule must come after all rules that match text keywords
            (r'[a-z_][a-z0-9_]*', Name),

            (r'((([0-9]*\.[0-9]+)|([0-9]+\.))([ed][-+]?[0-9]+)?)|([0-9]+[ed][-+]?[0-9]+)', Number.Float),
            (r'\d+', Number.Integer),

            (r'"[^"]*"', String.Double),

            (r'(-|\+|\*|\/|\^)', Operator),
            (r'(==|<=|>=|!=|<|>)', Operator),
            (r'(=|&|\||&&|\|\||!|:)', Operator),

            (r'[\[\]().,;]', Punctuation),

            (r'\s+', Whitespace)
        ],

        "macro_directive" : [
            # Do not pop state if there is a end-of-line continuation.
            # NB: this rule needs to consume the newline, so uses \n instead of $, otherwise
            # the next rule will pop the state.
            # Unfortunately the optional comment will not be highlighted, but there’s no easy fix.
            (r"\\\\[ \t]*(//.*)?\n", Text),
            # The following rule must come before the include(), otherwise it’s overriden by the \s+ rule
            (r"$", Whitespace, "#pop"),

            (words(dynare_macro_commands, prefix=r'\b', suffix=r'\b'), Keyword.Reserved),
            include("macro_common")
        ],

        "macro_expansion" : [
            (r"}", Punctuation, "#pop"),
            include("macro_common")
        ],


    }
