import json
import os
from copy import deepcopy
from pathlib import Path

import numpy as np

from chinatravel.data.load_datasets import load_query
from chinatravel.evaluation.commonsense_constraint import (
    evaluate_commonsense_constraints,
)
from chinatravel.evaluation.hard_constraint import evaluate_hard_constraints_v2
from chinatravel.evaluation.schema_constraint import evaluate_schema_constraints
from chinatravel.evaluation.utils import load_json_file
from chinatravel.symbol_verification.concept_func import func_dict
from chinatravel.symbol_verification.dsl import execute_dsl_code

EVALUATOR_REVISION = "LAMDA-NeSy/ChinaTravel@933744f"
SANDBOX_REVISION = "LAMDA-NeSy/ChinaTravel-Sandbox@3cd0d8f"
PROJECT_ROOT = Path(__file__).resolve().parent

DEFAULT_ATTRACTION_PREFERENCE = """
attraction_count = 0
for activity in allactivities(plan):
    if activity_type(activity) == 'attraction':
        attraction_count += 1
result=attraction_count/(4*day_count(plan))
"""

DEFAULT_TRANSPORT_PREFERENCE = """
time_cost = 0
transport_count = 0
for activity in allactivities(plan):
    transports = activity_transports(activity)
    if transports!=[]:
        transport_count += 1
        time_cost += innercity_transport_time(transports)
if transport_count > 0:
    average_time_cost = time_cost / transport_count
    result= (-1/105) * average_time_cost + 8/7
else:
    result=0
"""

DEFAULT_RESTAURANT_PREFERENCE = """
res_count=0
for activity in allactivities(plan):
    if activity_type(activity) in ['breakfast', 'lunch', 'dinner']:
        res_count+=1
res_count=res_count/(day_count(plan))
result=res_count/3
"""

DEFAULT_PREFERENCES = (
    DEFAULT_ATTRACTION_PREFERENCE,
    DEFAULT_TRANSPORT_PREFERENCE,
    DEFAULT_RESTAURANT_PREFERENCE,
)


def load_result(result_dir, query_index):
    json_files = {
        os.path.splitext(filename)[0]: os.path.join(root, filename)
        for root, _, filenames in os.walk(result_dir)
        for filename in filenames
        if filename.endswith(".json")
    }
    plans = {}
    matched_uid = []
    unmatched_uid = []
    for query_id in query_index:
        result_file = json_files.get(query_id)
        if result_file is None:
            plans[query_id] = {}
            unmatched_uid.append(query_id)
            continue
        try:
            plans[query_id] = load_json_file(result_file)
            matched_uid.append(query_id)
        except (OSError, ValueError, TypeError):
            plans[query_id] = {}
            unmatched_uid.append(query_id)
    return plans, matched_uid, unmatched_uid


def _default_preference_scores(query_index, result_data, all_pass_id):
    all_pass_id = set(all_pass_id)
    all_scores = []
    for query_id in query_index:
        if query_id not in all_pass_id:
            all_scores.append(np.zeros(len(DEFAULT_PREFERENCES)))
            continue

        plan = result_data[query_id]
        scores = []
        for constraint in DEFAULT_PREFERENCES:
            variables = deepcopy(func_dict)
            variables["plan"] = plan
            try:
                execute_dsl_code(
                    constraint,
                    variables,
                    allowed_builtins={"set": set},
                )
                value = float(variables.get("result", 0.0))
                scores.append(max(0.0, min(1.0, value)))
            except Exception:  # noqa: BLE001 - invalid preference code scores zero
                scores.append(0.0)
        all_scores.append(np.asarray(scores))

    if not all_scores:
        return np.zeros(len(DEFAULT_PREFERENCES))
    return np.mean(all_scores, axis=0)


def _frame_records(frame):
    return json.loads(frame.to_json(orient="records", force_ascii=False))


def evaluate(args, _result=None):
    query_index, query_data = load_query(args)
    result_data, matched_uid, unmatched_uid = load_result(
        args.result_dir,
        query_index,
    )

    result = {
        "split": args.splits,
        "language": args.lang,
        "evaluator_revision": EVALUATOR_REVISION,
        "sandbox_revision": SANDBOX_REVISION,
        "matched_uid": matched_uid,
        "unmatched_uid": unmatched_uid,
    }

    schema = load_json_file(
        PROJECT_ROOT / "chinatravel/evaluation/output_schema.json"
    )
    yield {"stage": "schema", "progress": 0}
    schema_rate, schema_details, schema_pass_id = evaluate_schema_constraints(
        query_index,
        result_data,
        schema=schema,
    )
    result["DR"] = schema_rate
    yield {"stage": "schema", "progress": 100}

    yield {"stage": "commonsense", "progress": 0}
    (
        epr_macro,
        epr_micro,
        commonsense_details,
        commonsense_pass_id,
    ) = evaluate_commonsense_constraints(
        query_index,
        query_data,
        result_data,
        verbose=False,
        lang=args.lang,
    )
    result["EPR_micro"] = epr_micro
    result["EPR_macro"] = epr_macro
    yield {"stage": "commonsense", "progress": 100}

    yield {"stage": "logic", "progress": 0}
    (
        lpr_macro,
        lpr_micro,
        conditional_lpr_macro,
        conditional_lpr_micro,
        logic_details,
        logic_pass_id,
    ) = evaluate_hard_constraints_v2(
        query_index,
        query_data,
        result_data,
        env_pass_id=commonsense_pass_id,
        verbose=False,
        lang=args.lang,
    )
    result["LPR_micro"] = lpr_micro
    result["LPR_macro"] = lpr_macro
    result["C-LPR"] = conditional_lpr_micro
    result["C-LPR_macro"] = conditional_lpr_macro
    yield {"stage": "logic", "progress": 100}

    all_pass_id = sorted(
        set(schema_pass_id)
        & set(commonsense_pass_id)
        & set(logic_pass_id)
    )
    fpr = len(all_pass_id) / len(query_index) * 100
    dav, att, ddr = _default_preference_scores(
        query_index,
        result_data,
        all_pass_id,
    ) * 100
    overall = (
        0.1 * epr_micro
        + 0.1 * epr_macro
        + 0.25 * conditional_lpr_micro
        + 0.05 * dav
        + 0.05 * att
        + 0.05 * ddr
        + 0.4 * fpr
    )
    result.update(
        {
            "FPR": fpr,
            "DAV": dav,
            "ATT": att,
            "DDR": ddr,
            "overall": overall,
            "all_pass_uid": all_pass_id,
            "details": {
                "schema": _frame_records(schema_details),
                "commonsense": _frame_records(commonsense_details),
                "logic": _frame_records(logic_details),
            },
        }
    )
    yield {"stage": "final", "progress": 100, "result": result}