"""Annotation consistency analysis: detect labeling style variations.
===================================================================
Checks:
1. Box tightness distribution (multi-modal = different annotator styles)
2. Cross-camera annotation bias
3. Potential labeling errors (outlier boxes)
4. Train/val distribution comparison
Runs on CPU.
"""
import sys, os, json
import numpy as np
from collections import defaultdict
from tqdm import tqdm

PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)


def main():
    train_dir = 'Data/Detection_dataset/labels/train'
    val_dir = 'Data/Detection_dataset/labels/val'
    train_img = 'Data/Detection_dataset/images/train'
    val_img = 'Data/Detection_dataset/images/val'

    def analyze(dir_path, img_path, label):
        areas_px = []; ratios = []; box_counts = []
        cam_stats = defaultdict(list)
        tiny_boxes = []

        for f in tqdm(sorted(os.listdir(dir_path)), desc=label):
            if not f.endswith('.txt'): continue
            img_f = f.replace('.txt', '.jpg')
            img_exists = os.path.exists(os.path.join(img_path, img_f))
            cam = img_f.split('_2025')[0] if '_2025' in img_f else 'unknown'

            n = 0
            with open(os.path.join(dir_path, f)) as fh:
                for line in fh:
                    p = line.strip().split()
                    if len(p) < 5: continue
                    w, h = float(p[3]), float(p[4])
                    area = w * h
                    areas_px.append(area * 3200 * 1800)  # approximate pixel area
                    if h > 0: ratios.append((w*3200)/(h*1800))
                    cam_stats[cam].append({'w': w, 'h': h, 'area': area})
                    n += 1
                    # Flag potentially problematic boxes
                    if w < 0.003 or h < 0.003:  # absurdly tiny
                        tiny_boxes.append((f, w, h))
            box_counts.append(n)

        return {
            'areas': np.array(areas_px),
            'ratios': np.array(ratios),
            'box_counts': np.array(box_counts),
            'cam_stats': {cam: {'n': len(v), 'mean_w': float(np.mean([x['w'] for x in v])),
                                'std_w': float(np.std([x['w'] for x in v])),
                                'mean_h': float(np.mean([x['h'] for x in v])),
                                'std_h': float(np.std([x['h'] for x in v]))}
                         for cam, v in cam_stats.items()},
            'tiny_boxes': tiny_boxes,
            'n_labels': len(areas_px),
        }

    train_s = analyze(train_dir, train_img, 'Train')
    val_s = analyze(val_dir, val_img, 'Val')

    # Check for multimodal width/height distributions (indicates annotator style differences)
    from scipy import stats as scipy_stats
    try:
        train_w = train_s['areas'] ** 0.5
        kde = scipy_stats.gaussian_kde(train_w[:5000])  # sample for speed
        x = np.linspace(np.percentile(train_w, 1), np.percentile(train_w, 99), 100)
        peaks = []
        y = kde(x)
        for i in range(1, len(y)-1):
            if y[i] > y[i-1] and y[i] > y[i+1] and y[i] > y.max()*0.3:
                peaks.append(float(x[i]))
        multimodal = len(peaks) > 1
    except:
        multimodal = False
        peaks = []

    print('='*60)
    print('ANNOTATION QUALITY REPORT')
    print('='*60)
    print(f'\nTrain labels: {train_s["n_labels"]} boxes in {len(train_s["box_counts"])} files')
    print(f'Val labels:   {val_s["n_labels"]} boxes in {len(val_s["box_counts"])} files')
    print(f'Median boxes/img — Train: {np.median(train_s["box_counts"]):.0f}  Val: {np.median(val_s["box_counts"]):.0f}')

    # Per-camera annotation consistency
    print('\nPer-camera box statistics (Train):')
    print(f'{"Camera":12s} {"Boxes":>6s} {"mean W":>8s} {"std W":>8s} {"mean H":>8s} {"std H":>8s}')
    print('-'*55)
    for cam in sorted(train_s['cam_stats']):
        s = train_s['cam_stats'][cam]
        print(f'{cam:12s} {s["n"]:6d} {s["mean_w"]:8.4f} {s["std_w"]:8.4f} {s["mean_h"]:8.4f} {s["std_h"]:8.4f}')

    # Outlier detection
    print(f'\nTrain/Val distribution comparison:')
    print(f'  Train median size: {np.median(train_s["areas"]):.0f} px²')
    print(f'  Val median size:   {np.median(val_s["areas"]):.0f} px²')
    size_diff = np.median(val_s['areas'])/np.median(train_s['areas']) - 1
    print(f'  Difference: {size_diff*100:+.1f}%')
    if abs(size_diff) > 0.05:
        print(f'  ⚠ Distribution mismatch > 5% — val may not represent train!')

    if multimodal:
        print(f'\nMultimodal size distribution detected! ({len(peaks)} peaks)')
        print(f'  Possible multiple annotator styles. Peak sizes: {[f"{p:.0f}" for p in peaks]} px')
    else:
        print(f'\nSize distribution appears unimodal — consistent annotation style.')

    if train_s['tiny_boxes']:
        print(f'\n⚠ {len(train_s["tiny_boxes"])} suspiciously tiny boxes found:')
        for f, w, h in train_s['tiny_boxes'][:5]:
            print(f'  {f}: w={w:.4f}, h={h:.4f}')

    with open('logs/annotation_quality.json', 'w') as f:
        json.dump({
            'multimodal': multimodal,
            'peaks': peaks,
            'train_val_diff': round(size_diff, 4),
            'cam_stats': {cam: {k: round(v, 4) if isinstance(v, float) else v for k, v in s.items()}
                         for cam, s in train_s['cam_stats'].items()},
        }, f, indent=2)
    print('\nSaved to logs/annotation_quality.json')


if __name__ == '__main__':
    main()