"""
Plot correlation between GuideScan2 specificity scores (new pipeline) and
MIT specificity scores (old CRISPRDesigner pipeline).

Usage:
    python workflow/scripts/plot_specificity_correlation.py \
        --new_guides results/GuideDesign/designGuides.txt \
        --reference tests/data/GATA1PPIF_reference_specificity.tsv \
        --output results/specificity_score_correlation.png
"""

import argparse
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.stats import spearmanr


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--new_guides', required=True,
                        help='designGuides.txt from new pipeline')
    parser.add_argument('--reference', required=True,
                        help='Reference specificity TSV with specificity_score_MIT column')
    parser.add_argument('--output', required=True,
                        help='Output PNG path')
    args = parser.parse_args()

    new = pd.read_csv(args.new_guides, sep='\t')
    ref = pd.read_csv(args.reference, sep='\t')

    new['protospacer'] = new['GuideSequenceWithPAM'].str[:20]
    merged = new[['protospacer', 'specificity_score']].merge(
        ref[['protospacer', 'specificity_score_MIT']], on='protospacer', how='inner'
    )

    # Replace NaN with 0
    merged['specificity_score'] = merged['specificity_score'].fillna(0)
    merged['specificity_score_MIT'] = merged['specificity_score_MIT'].fillna(0)

    r, p = spearmanr(merged['specificity_score'], merged['specificity_score_MIT'])

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.scatter(merged['specificity_score_MIT'], merged['specificity_score'],
               alpha=0.6, edgecolors='none', s=40, color='steelblue')

    ax.set_xlim(0, 100)
    ax.set_ylim(0, 1)
    ax.set_xlabel('MIT specificity score (old pipeline, hg19)', fontsize=13)
    ax.set_ylabel('GuideScan2 specificity score (new pipeline, hg38)', fontsize=13)
    ax.set_title(
        f'Off-target specificity: GuideScan2 vs MIT\n'
        f'Spearman r = {r:.3f}, p = {p:.1e}, n = {len(merged)}',
        fontsize=12
    )
    ax.text(2, 0.96, 'NaN replaced with 0', fontsize=9, color='gray', va='top')

    plt.tight_layout()
    plt.savefig(args.output, dpi=150)
    print(f'Saved: {args.output}')


if __name__ == '__main__':
    main()
