# Copyright (c) 2024-present, Royal Bank of Canada.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import logging
logging.disable(logging.CRITICAL)
import copy
import os
import numpy as np
import random
from absl import app
from ml_collections.config_dict import ConfigDict
from ml_collections import config_flags
import uuid
import logging
from gpt_client import GPTClient, GPTConfig
from domains import Domain, PDDLEnv
from planning import evaluate_action_level_planning, PlanningStrategy, evaluate_planning_on_problem_candidates, \
    evaluate_all_tasks, generate_one_domain_file, rate_and_update_files
import coloredlogs
import json
import wandb
import pdb
import prompts

from problem_domain_translation import translate_problems_given_one_task, \
    generate_exact_n_problem_translation_candidates
from pddl_utils import PDDLObj, validate_problem_pddl
from utils import extract_code, get_function_from_code, harmonic_mean

_CONFIG = config_flags.DEFINE_config_dict(
    'cfg',
    ConfigDict(dict(
        run_id=str(uuid.uuid4())[:8],
        debug=False,
        data_path_example='./data_my/domains',
        data_path='./data_my/domains/test',
        exp_path='./experiments',
        source_folder='gpt4-old-top1-trans1_seq_rowcolumn_dir_eva5',
        log_prefix="test",
        seed=42,
        gpt_args=dict(
            api_key="your-openai-api-key",
            model_name='gpt-4o-2024-08-06',
            multimodal=True
        ),
        env_args=dict(
            fd_py_path='/path/to/downward/fast-downward.py',
            fd_search_time_limit=300,
            val_bin_path='/path/to/VAL/build/linux64/Release/bin/Validate',
        ),
        planning_strategy_args=dict(
            turns=4,  # How many turns to use for the conversation with LLM
            best_of_n=1,  # How many samples to generate from LLM and choose the best one
            vlm_path='',
            multimodal=True,
            sequence=True
        ),
        problem_translation_args=dict(
            active=True,  # Whether to generate problem translation candidates, or use the target problem
            n_candidates=5,  # Number of candidates to generate
            one_domain_per_candidate=True,
            logprob_selection=False,
            # Whether to use logprob for selecting the best candidate among logprob_candidates candidates
            logprob_candidates=0,
            add_domain_proposal=False,  # Whether to add a domain proposal before the problem translation candidates
            vlm_path='',
        ),
        first_task_index=0,
        max_tasks=100,  # Maximum number of tasks to evaluate
        context_domain_name='maze', # This is strict, all the prompts are based on blocksworld
        target_domain_name='frozenlake',
        wandb_args=dict(
            project="llm-planning",
            entity="your-wandb-entity",
        ),
        exp_flags=dict(
        ),
    ))
)


def run(cfg):
    
    # logging.getLogger().handlers.clear()
    # coloredlogs.install(level="INFO")
    # logging.info(f"Running with config: {cfg.to_dict()}")
    # assert cfg.context_domain_name == 'blocksworld', "Context domain must be blocksworld"
    assert cfg.context_domain_name != cfg.target_domain_name, "Context and target domains must be different."
    _seed_all(cfg.seed)
    run_exp_dir = os.path.join(cfg.exp_path, f"{cfg.log_prefix}/seed_{cfg.seed}")
    os.makedirs(run_exp_dir, exist_ok=True)
    summary_log_path = os.path.join(run_exp_dir, "summary_logs.json")
    if os.path.exists(summary_log_path):
        logging.warning(
            f"Summary log path already exists: {summary_log_path}, either delete it or change the prefix or seed"
        )
        exit(0)
    # wandb_run = _config_wandb(cfg)

    gpt_client = None#GPTClient(GPTConfig(**cfg.gpt_args))
    # initialize two domains, these require the path contain pddl files
    context_domain = Domain(cfg.data_path, cfg.context_domain_name, vlm_path = cfg.problem_translation_args.vlm_path)
    target_domain_example = Domain(cfg.data_path_example, cfg.target_domain_name, vlm_path = cfg.problem_translation_args.vlm_path)
    target_domain = Domain(cfg.data_path, cfg.target_domain_name, vlm_path = cfg.problem_translation_args.vlm_path)
    # pdb.set_trace()
    # initialize a pddl env, can search/validate plan
    pddl_env = PDDLEnv(**cfg.env_args)
    first_task_index = cfg.first_task_index
    source_folder=cfg.source_folder
    
    source_folder = f'./experiments/paper_table_2/{cfg.context_domain_name}/{cfg.target_domain_name}/{source_folder}/seed_{first_task_index}'
    print(source_folder)
    try:
        with open(os.path.join(source_folder, 'domain.txt'), 'r') as f:
            best_generated_pddl = f.read()
            print(best_generated_pddl)
        
        with open(os.path.join(source_folder, 'problem.txt'), 'r') as f:
            best_generated_problem_pddl = f.read()
            print(best_generated_problem_pddl)
    except:
        exit(0)
    
    # pdb.set_trace()
    if 'modify_action' in best_generated_pddl:
        # pdb.set_trace()
        target_domain_template_pddl = target_domain_example.get_domain_template_pddl()
        pddl_obj = PDDLObj.from_pddl_str(target_domain_template_pddl, domain_pddl_template=target_domain_template_pddl)
        pddl_obj.modify_domain(best_generated_pddl)
        best_generated_pddl = pddl_obj.to_str()
        
    target_problem_list = [target_domain.get_task_pddl(i) for i in range(cfg.max_tasks)]
    other_task_nls = [target_domain.get_task_nl(i) for i in range(cfg.max_tasks)]
    other_task_templates = [target_domain.get_task_template(i) for i in range(cfg.max_tasks)]
    other_task_image_paths = [target_domain.get_task_image(i) for i in range(cfg.max_tasks)]
    other_task_images = [prompts.encode_image(target_domain.get_task_image(i)) for i in range(cfg.max_tasks)]
    gen_problem_list = translate_problems_given_one_task(
        gpt_client=gpt_client, domain_pddl=best_generated_pddl, domain_nl=target_domain.get_domain_nl(),
        context_problem_pddl=best_generated_problem_pddl,
        context_problem_nl=target_domain_example.get_task_nl(first_task_index),
        context_problem_template_pddl=target_domain_example.get_task_template(first_task_index),
        target_problem_nls=other_task_nls,
        target_problem_templates=other_task_templates,
        target_task_images=other_task_images,
        save_folder = run_exp_dir
    )
    
    correct_tasks_frac, valid = evaluate_all_tasks(
        pddl_env=pddl_env,
        target_domain_pddl=target_domain_example.get_domain_pddl(),
        target_domain_problem_pddls=target_problem_list,
        target_gen_domain_pddl=best_generated_pddl,
        target_gen_problem_pddls=gen_problem_list,
        other_task_image_paths=other_task_image_paths,
        vlm_path = cfg.planning_strategy_args.vlm_path,
        multimodal = cfg.planning_strategy_args.multimodal,
        sequence = cfg.planning_strategy_args.sequence,
        exp_flags=cfg.exp_flags,
    )
    
    summary_metrics = {
        'correct_tasks_score': correct_tasks_frac,
        'valid': valid,
        # 'rw_t_to_gen_frac': rw_t_to_gen_frac,
        # 'rw_gen_to_t_frac': rw_gen_to_t_frac,
        # 'used_prompt_tokens': gpt_client.used_prompt_tokens,
        # 'used_completion_tokens': gpt_client.used_completion_tokens,
        # 'cost_dollars': gpt_client.get_cost(),
    }
    # wandb_run.summary.update(summary_metrics)
    # summary_metrics = {}
    # gen_problem_list = []
    
    # gpt_client.save_chats(save_dir=os.path.join(run_exp_dir, "chats"))
    # file_logger = get_file_logger(os.path.join(run_exp_dir, f"run.log"))
    # Save aux as file
    summary_log_dict = {
        # 'aux': aux,
        'cfg': cfg.to_dict(),
        'summary_metrics': summary_metrics,
        'gen_problem_list': gen_problem_list,
        'task_0_problem_translation_candidates': best_generated_problem_pddl,
        'best_gen_domain_pddl': best_generated_pddl,

    }
    with open(summary_log_path, 'w') as f:
        json.dump(summary_log_dict, f, indent=2)
    # wandb_run.save(summary_log_path)
    # file_logger.info(f"Used tokens: {gpt_client.used_tokens}")



def main(_):
    cfg = _CONFIG.value
    run(cfg)


def get_file_logger(log_path):
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    fh = logging.FileHandler(log_path)
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)
    return logger


def _seed_all(seed):
    """Sets the random seed"""
    np.random.seed(seed)
    random.seed(seed)


def _config_wandb(cfg):
    log_prefix = cfg.log_prefix
    debug = cfg.debug
    wandb_run = wandb.init(
        **cfg.wandb_args, mode="disabled" if debug else None,
        config=cfg.to_dict(), name=f"{log_prefix}/{cfg.run_id}"
    )
    wandb.config.update(cfg.to_dict())
    # wandb.run.save()
    return wandb_run


if __name__ == '__main__':
    app.run(main)
