#!/usr/bin/env python3
"""
Build pair index from 3DMatch evaluation logs
"""
import json
import sys
import os
from pathlib import Path
from collections import defaultdict

def parse_gt_log(log_file):
    """Parse gt.log evaluation file
    
    Format:
    source_id target_id ...
    r11 r12 r13 tx
    r21 r22 r23 ty
    r31 r32 r33 tz
    0   0   0   1
    """
    pairs = []
    
    with open(log_file, 'r') as f:
        lines = [line.strip() for line in f.readlines() if line.strip()]
    
    i = 0
    while i < len(lines):
        # Parse header line
        header = lines[i].split()
        if len(header) < 2:
            i += 1
            continue
        
        try:
            source_id = int(header[0])
            target_id = int(header[1])
        except (ValueError, IndexError):
            i += 1
            continue
        
        # Parse 4x4 transformation matrix
        if i + 4 >= len(lines):
            break
        
        try:
            matrix = []
            for j in range(1, 5):
                row = [float(x) for x in lines[i + j].split()]
                if len(row) != 4:
                    raise ValueError(f"Invalid row: {lines[i + j]}")
                matrix.append(row)
            
            pairs.append({
                'source_id': source_id,
                'target_id': target_id,
                'gt_transform': matrix
            })
            i += 5
        except (ValueError, IndexError) as e:
            print(f"Warning: Failed to parse transformation matrix at line {i}: {e}")
            i += 1
    
    return pairs

def find_ply_file(scene_dir, fragment_id):
    """Find cloud_bin_<id>.ply file"""
    pattern = f"cloud_bin_{fragment_id}.ply"
    for root, dirs, files in os.walk(scene_dir):
        if pattern in files:
            return str(Path(root) / pattern)
    return None

def build_pair_index(raw_root, output_file):
    """Build pair index for redkitchen scene"""
    raw_root = Path(raw_root)
    pairs = []
    
    scene_id = '7-scenes-redkitchen'
    scene_dir = raw_root / scene_id
    eval_dir = raw_root / f'{scene_id}-evaluation'
    
    print(f"Processing: {scene_id}")
    
    if not scene_dir.exists():
        print(f"Error: Scene folder not found: {scene_dir}")
        return False
    
    if not eval_dir.exists():
        print(f"Error: Evaluation folder not found: {eval_dir}")
        return False
    
    # Find gt.log or any .log file
    log_file = None
    if (eval_dir / 'gt.log').exists():
        log_file = eval_dir / 'gt.log'
    else:
        for f in eval_dir.rglob('*.log'):
            log_file = f
            break
    
    if not log_file:
        print(f"Error: No .log file found in {eval_dir}")
        return False
    
    print(f"Parsing: {log_file.relative_to(raw_root)}")
    scene_pairs = parse_gt_log(log_file)
    
    # Validate and add source/target paths
    added = 0
    for pair in scene_pairs:
        source_file = find_ply_file(scene_dir, pair['source_id'])
        target_file = find_ply_file(scene_dir, pair['target_id'])
        
        if source_file and target_file and pair['source_id'] != pair['target_id']:
            pair['scene_id'] = scene_id
            pair['source_path'] = source_file
            pair['target_path'] = target_file
            pairs.append(pair)
            added += 1
        else:
            if not source_file:
                print(f"  Warning: Source file not found for id {pair['source_id']}")
            if not target_file:
                print(f"  Warning: Target file not found for id {pair['target_id']}")
    
    print(f"Added {added} pairs from {len(scene_pairs)} total in log")
    
    # Save index
    output_file = Path(output_file)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    
    with open(output_file, 'w') as f:
        json.dump(pairs, f, indent=2)
    
    print(f"\nPair index saved: {output_file}")
    print(f"Total pairs: {len(pairs)}")
    return len(pairs) > 0

if __name__ == '__main__':
    raw_root = sys.argv[1] if len(sys.argv) > 1 else 'data/raw/3dmatch'
    output_file = sys.argv[2] if len(sys.argv) > 2 else 'data/processed/pair_index.json'
    
    success = build_pair_index(raw_root, output_file)
    sys.exit(0 if success else 1)