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

For each metric range, it finds the configuration with the lowest query time,
provided that its metric 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), metric range bounds (min/max), and step size
  - Metric column name (default: 'Recall', but can be any metric like 'RR@10', 'Success@5', etc.)

Output:
  - A TSV file with one row per metric 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 (Recall - default):
  python scripts/find_best_grid_results_by_recall_range.py \
      path/to/final_report.tsv \
      path/to/best_results.tsv \
      --metric-min 0.6 --metric-max 0.99 --step 0.01

Example usage (IR metric like RR@10):
  python scripts/find_best_grid_results_by_recall_range.py \
      path/to/final_report.tsv \
      path/to/best_results.tsv \
      --metric-col 'RR@10' \
      --metric-min 0.3 --metric-max 0.5 --step 0.02
"""

import pandas as pd
import argparse
import os

def find_best_by_metric_range(df, metric_col='Recall', space_limit=None, metric_min=0.90, metric_max=0.995, step=0.005):
    ranges = []
    start = metric_min
    while start < metric_max:
        end = min(start + step, metric_max)
        ranges.append((start, end))
        start = end

    selected_rows = []
    to_print = []

    target_col_name = f'Target {metric_col}'

    for mmin, mmax in ranges:
        subset = df[(df[metric_col] >= mmin) & (df[metric_col] < mmax)]
        print(f"🔍 Processing {metric_col} range: [{mmin:.3f}, {mmax:.3f}) with {len(subset)} candidates")
        if space_limit is not None:
            subset = subset[subset['Memory Usage (Bytes)'] <= space_limit]
        if not subset.empty:
            best_row = subset.loc[subset['Query Time (microsecs)'].idxmin()].copy()
            best_row[target_col_name] = mmin
            selected_rows.append(best_row)
            metric_val = best_row[metric_col]
            to_print.append(f"{metric_col}={metric_val:.4f}, Query Time={best_row['Query Time (microsecs)']}μs, Memory={best_row['Memory Usage (Bytes)']}B")
        

    result_df = pd.DataFrame(selected_rows)
    if not result_df.empty:
        cols = [target_col_name] + [c for c in result_df.columns if c != target_col_name]
        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 metric 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("--metric-col", type=str, default="Recall", help="Metric column name to partition by (default: Recall). Can be any metric like 'RR@10', 'Success@5', etc.")
    parser.add_argument("--space_usage", type=int, default=None, help="Optional space usage limit in bytes")
    parser.add_argument("--metric-min", type=float, default=0.900, help="Minimum metric value for range partition")
    parser.add_argument("--metric-max", type=float, default=0.995, help="Maximum metric value for range partition")
    parser.add_argument("--step", type=float, default=0.005, help="Step size for metric ranges")

    args = parser.parse_args()
    metric_col = args.metric_col

    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())
    
    # Handle legacy Accuracy -> Recall renaming
    if metric_col == "Recall" and 'Accuracy' in df.columns and 'Recall' not in df.columns:
        df = df.rename(columns={'Accuracy': 'Recall'})
        print("✅ Renamed 'Accuracy' column to 'Recall'")

    # Check that metric column exists
    if metric_col not in df.columns:
        print(f"Error: Metric column '{metric_col}' not found in input file.")
        print(f"Available columns: {', '.join(df.columns)}")
        return

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

    print(f"{metric_col} range in dataset: {df[metric_col].min()} to {df[metric_col].max()}")

    min_metric_in_df = df[metric_col].min()

    # Auto-adjust percentage ranges if metric is in 0-100 range
    if min_metric_in_df > 1.0 and args.metric_min <= 1.0:
        print(f"Metric expressed in percentage (0-100). Adjusting metric_min from {args.metric_min} to {args.metric_min * 100}, metric_max from {args.metric_max} to {args.metric_max * 100}, step from {args.step} to {args.step * 100}")
        args.metric_min *= 100
        args.metric_max *= 100
        args.step *= 100
        
    best_df = find_best_by_metric_range(
        df,
        metric_col=metric_col,
        space_limit=args.space_usage,
        metric_min=args.metric_min,
        metric_max=args.metric_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()
