import os
import _init_paths
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = [8, 8]

from lib.test.analysis.plot_results import print_results
from lib.test.evaluation import get_dataset, trackerlist

PROJECT_DIR = os.environ.get(
    'ATCTRACK_VLM_PROJECT_DIR',
    os.path.abspath(os.path.join(os.path.dirname(__file__), '..')),
)
SAVE_DIR = os.environ.get('SAVE_DIR', os.path.join(PROJECT_DIR, 'output'))
DATASET_ROOT = os.environ.get('DATASET_ROOT', os.path.abspath(os.path.join(PROJECT_DIR, '../Dataset')))

os.environ.setdefault('ATCTRACK_VLM_PROJECT_DIR', PROJECT_DIR)
os.environ.setdefault('SAVE_DIR', SAVE_DIR)
os.environ.setdefault('DATASET_ROOT', DATASET_ROOT)
os.environ.setdefault('LASOT_LANG_DIR', os.path.join(DATASET_ROOT, 'LaSOT/LaSOTBenchmark'))

dataset_name = os.environ.get('ANALYSIS_DATASET', 'lasot_lang')
tracker_param = os.environ.get('ANALYSIS_TRACKER_PARAM', 'atctrack_qwen_state')
result_dir = os.environ.get(
    'ANALYSIS_RESULT_DIR',
    os.path.join(SAVE_DIR, tracker_param, dataset_name),
)

trackers = trackerlist(
    name='atctrack',
    parameter_name=tracker_param,
    dataset_name=dataset_name,
    run_ids=None,
    display_name='ATCTrack-QwenState-9B',
)

# test.py writes this run to output/<tracker_param>/<dataset_name>, while the
# default analysis path is output/test/tracking_results/<tracker>/<param>.
trackers[0].results_dir = result_dir
trackers[0].results_eval_dir = result_dir

print(f'dataset_name: {dataset_name}')
print(f'tracker_param: {tracker_param}')
print(f'result_dir: {result_dir}')

dataset = get_dataset(dataset_name)

print_results(
    trackers,
    dataset,
    dataset_name,
    merge_results=True,
    plot_types=('success', 'prec', 'norm_prec'),
    force_evaluation=True,
)