"""
Implement different utilities with OptiMUS v0.3
"""
import os
import sys
from typing import Dict

from constraint import get_constraints
from constraint_model import get_constraint_formulations
from objective import get_objective
from objective_model import get_objective_formulation
from parameters import get_params
from target_code import get_codes
from utils import load_state, save_state, Logger

sys.path.append(os.path.join(os.path.dirname(__file__), "."))

ERROR_CORRECTION = False
# MODEL = "gpt-4o"
MODEL = "glm-4-0520"
RAG_MODE = None
DEFAULT_LABELS = {"types": ["Mixed Integer Linear Programming"], "domains": ["Energy Optimization"]}


def get_intro_latex_code_map(fname) -> Dict:
    """
    Extract the internal OptiMUS data structure that maps natural language, LaTeX and code
    No data is required by this utility
    """
    run_dir = "."

    # Extract parameters from the natural language description
    with open(fname, "r") as f:
        desc = f.read()
        f.close()

    params = get_params(desc, check=True)

    # Construct the initial state
    state = {"description": desc, "parameters": params}
    save_state(state, "state_1_params.json")

    logger = Logger(f"log.txt")
    logger.reset()

    # Get objective
    state = load_state(os.path.join(run_dir, "state_1_params.json"))
    objective = get_objective(
        state["description"],
        state["parameters"],
        check=ERROR_CORRECTION,
        logger=logger,
        model=MODEL,
        rag_mode=RAG_MODE,
        labels=DEFAULT_LABELS,
    )
    print(objective)
    state["objective"] = objective
    save_state(state, os.path.join(run_dir, "state_2_objective.json"))

    # Get constraints
    state = load_state(os.path.join(run_dir, "state_2_objective.json"))
    constraints = get_constraints(
        state["description"],
        state["parameters"],
        check=ERROR_CORRECTION,
        logger=logger,
        model=MODEL,
        rag_mode=RAG_MODE,
        labels=DEFAULT_LABELS,
    )
    print(constraints)
    state["constraints"] = constraints
    save_state(state, os.path.join(run_dir, "state_3_constraints.json"))

    # Get constraint formulations
    state = load_state(os.path.join(run_dir, "state_3_constraints.json"))
    constraints, variables = get_constraint_formulations(
        state["description"],
        state["parameters"],
        state["constraints"],
        check=ERROR_CORRECTION,
        logger=logger,
        model=MODEL,
        rag_mode=RAG_MODE,
        labels=DEFAULT_LABELS,
    )
    state["constraints"] = constraints
    state["variables"] = variables
    save_state(state, os.path.join(run_dir, "state_4_constraints_modeled.json"))

    # Get objective formulation
    state = load_state(os.path.join(run_dir, "state_4_constraints_modeled.json"))
    objective = get_objective_formulation(
        state["description"],
        state["parameters"],
        state["variables"],
        state["objective"],
        model=MODEL,
        check=ERROR_CORRECTION,
        rag_mode=RAG_MODE,
        labels=DEFAULT_LABELS,
    )
    state["objective"] = objective
    print("DONE OBJECTIVE FORMULATION")
    save_state(state, os.path.join(run_dir, "state_5_objective_modeled.json"))

    # # ####### Get codes
    state = load_state(os.path.join(run_dir, "state_5_objective_modeled.json"))
    constraints, objective = get_codes(
        state["description"],
        state["parameters"],
        state["variables"],
        state["constraints"],
        state["objective"],
        model=MODEL,
        check=ERROR_CORRECTION,
    )
    state["constraints"] = constraints
    state["objective"] = objective
    save_state(state, os.path.join(run_dir, "state_6_code.json"))

    return state


if __name__ == "__main__":
    get_intro_latex_code_map("./data/desc.txt")