"""
This script selects the best configurations from a grid search result file (final_report.tsv),
based on recall ranges and an optional memory usage constraint.

For each recall range, it finds the configuration with the lowest query time,
provided that its recall is within the range and (optionally) its memory usage
is below a specified threshold.

Inputs:
  - A final_report.tsv file generated by scripts/gather_grid_search_results.py
  - Optional filters: space usage (in bytes), recall range bounds, and step size

Output:
  - A TSV file with one row per recall subrange, showing the best configuration
    (i.e., lowest query time among those satisfying the constraints)

Workflow:
  1. First, run a grid search using:
        scripts/run_grid_search.py
  2. Then, collect results using:
        scripts/gather_grid_search_results.py
  3. Finally, filter the best results using this script.

Example usage:
  python scripts/find_best_grid_results_by_recall_range.py \
      path/to/final_report.tsv \
      path/to/best_results.tsv \
      --space_usage 50000000 \
      --recall_min 0.6 \
      --recall_max 0.99 \
      --step 0.01
"""

import pandas as pd
import argparse
import os

def find_best_by_recall_range(df, space_limit=None, recall_min=0.90, recall_max=0.995, step=0.005):
    ranges = []
    start = recall_min
    while start < recall_max:
        end = min(start + step, recall_max)
        ranges.append((start, end))
        start = end

    selected_rows = []

    to_print = []

    for rmin, rmax in ranges:
        subset = df[(df['Recall'] >= rmin) & (df['Recall'] < rmax)]
        print(f"🔍 Processing recall range: [{rmin:.3f}, {rmax:.3f}) with {len(subset)} candidates")
        if space_limit is not None:
            subset = subset[subset['Memory Usage (Bytes)'] <= space_limit]
        # subset = subset[subset['query.early-termination'] == "none"]
        if not subset.empty:
            best_row = subset.loc[subset['Query Time (microsecs)'].idxmin()].copy()
            best_row['Target Recall'] = rmin
            selected_rows.append(best_row)
            to_print.append(f"Recall={best_row['Recall']:.4f}, Query Time={best_row['Query Time (microsecs)']}μs, Memory={best_row['Memory Usage (Bytes)']}B, ef_search: {best_row.get('query.ef-search', 'N/A')}, patience: {best_row.get('query.patience', 'N/A')}, proximity-threshold: {best_row.get('query.proximity-threshold', 'N/A')}")
        

    result_df = pd.DataFrame(selected_rows)
    if not result_df.empty:
        cols = ['Target Recall'] + [c for c in result_df.columns if c != 'Target Recall']
        result_df = result_df[cols]

    for line in to_print:
        print(line)

    return result_df

def main():
    parser = argparse.ArgumentParser(description="Find best grid search results by recall range")
    parser.add_argument("input_tsv", type=str, help="Path to input final_report.tsv")
    parser.add_argument("output_tsv", type=str, help="Path to output best_results.tsv")
    parser.add_argument("--space_usage", type=int, default=None, help="Optional space usage limit in bytes")
    parser.add_argument("--recall_min", type=float, default=0.900, help="Minimum recall for range partition")
    parser.add_argument("--recall_max", type=float, default=0.995, help="Maximum recall for range partition")
    parser.add_argument("--step", type=float, default=0.005, help="Step size for recall ranges")

    args = parser.parse_args()

    if not os.path.exists(args.input_tsv):
        print(f"Error: Input file {args.input_tsv} does not exist.")
        return

    df = pd.read_csv(args.input_tsv, sep="\t")
    print(f"✅ Loaded {len(df)} rows from {args.input_tsv}")
    print(df.head())
    if 'Accuracy' in df.columns and 'Recall' not in df.columns:
        df = df.rename(columns={'Accuracy': 'Recall'})
        print("✅ Renamed 'Accuracy' column to 'Recall'")

    print("Recall range in dataset:", df["Recall"].min(), "to", df["Recall"].max())

    min_recall_in_df = df["Recall"].min()

    required_cols = {'Recall', 'Query Time (microsecs)', 'Memory Usage (Bytes)'}
    if not required_cols.issubset(set(df.columns)):
        print("Error: Input file must contain Recall, Query Time (microsecs), and Memory Usage (Bytes) columns.")
        return

    if min_recall_in_df > 1.0 and args.recall_min <= 1.0:
        print(f"Recall expressend in percentage (0-100). Adjusting recall_min from {args.recall_min} to {args.recall_min * 100}, recall_max from {args.recall_max} to {args.recall_max * 100}, step from {args.step} to {args.step * 100}")
        args.recall_min *= 100
        args.recall_max *= 100
        args.step *= 100
        
    best_df = find_best_by_recall_range(
        df,
        space_limit=args.space_usage,
        recall_min=args.recall_min,
        recall_max=args.recall_max,
        step=args.step,
    )

    best_df.to_csv(args.output_tsv, sep="\t", index=False)
    print(f"✅ Best results saved to {args.output_tsv}")

if __name__ == "__main__":
    main()
