# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import collections
import glob
import itertools
import json
import operator
import os
import re

import _jsonnet
import IPython.display
import natsort
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

# #%matplotlib inline
#plt.rcParams['figure.figsize'] = [20, 20]
sns.set(color_codes=True)

from seq2struct.utils import evaluation
from seq2struct.utils import registry
from seq2struct import datasets
# -

os.chdir('..')


# # Procedure
# 1. Identify all checkpoints that we had selected
#   - Snapshots that maximize sentence BLEU, exact match
#   - Idiom models: across all contains-hole/none, cov-examples/cov-xent, 10/20/40/80
#   - Baseline: across all evaluated steps for one model
# 2. Determine when and where idioms were output while decoding
# 3. Compute teacher-forced precision/recall
# 4. When the baseline gets it wrong and the model gets it right

# # Finding checkpoints for idiom models

def find_idiom_checkpoints():
    accuracy_per_run = collections.defaultdict(dict)
    all_metrics = []
    metric_types = set()
    rows = []
    for d in sorted(glob.glob('logdirs/20190201-hs-allmatches-anysplit-multimean/*')):
        exp_name = os.path.basename(d)
        exp_vars = re.match('filt-([^_]+)_st-([^_]+)_nt-([^_]+)', exp_name).groups()

        infer_paths = glob.glob(os.path.join(d, 'infer-val-step*-bs1.jsonl'))
        all_scores = []
        for infer_path in infer_paths:
            step = int(re.search('step(\d+)', infer_path).group(1))
            _, metrics = evaluation.compute_metrics(
                'configs/hearthstone/nl2code.jsonnet', '', 'val', infer_path)
            all_scores.append((step, metrics['exact match']))
            all_metrics.append((exp_name, step, metrics))
            metric_types.update(metrics.keys())

        all_scores.sort(key=operator.itemgetter(0))
        sorted_scores = sorted(all_scores, reverse=True, key=operator.itemgetter(1))
        rows.append(exp_vars + (len(all_scores),) + (sorted_scores[0] if sorted_scores else (-1, -1)))
        accuracy_per_run[exp_name] = {
            'x': [s[0] for s in all_scores],
            'all': [s[1] for s in all_scores],
        }
        print(d)

    metric_types = tuple(sorted(metric_types))
    df = pd.DataFrame(rows, columns=('filt', 'cov', 'nt', 'num steps eval', 'step', 'exact match'))
    flat_df = pd.DataFrame(
        [(exp_name, step) + tuple(metrics.get(t) for t in metric_types) for exp_name, step, metrics in all_metrics],
        columns=('exp_name', 'step') + metric_types)
    
    return flat_df


idiom_df = find_idiom_checkpoints()
print(idiom_df.loc[idiom_df['corpus BLEU'].idxmax()])
print(idiom_df.loc[idiom_df['exact match'].idxmax()])


# # Finding checkpoints for baseline models

def find_baseline_checkpoints():
    accuracy_per_run = collections.defaultdict(dict)
    all_metrics = []
    metric_types = set()
    rows = []
    for d in sorted(glob.glob('logdirs/20181231-nl2code-hearthstone-fef2c5b//*')):
        exp_name = os.path.basename(d)
        exp_vars = re.match('att([^_]+)', exp_name).groups()

        infer_paths = glob.glob(os.path.join(d, 'infer-val-step*-bs1.jsonl'))
        all_scores = []
        for infer_path in infer_paths:
            step = int(re.search('step(\d+)', infer_path).group(1))
            _, metrics = evaluation.compute_metrics(
                'configs/hearthstone/nl2code.jsonnet', '', 'val', infer_path)
            all_scores.append((step, metrics['exact match']))
            all_metrics.append((exp_name, step, metrics))
            metric_types.update(metrics.keys())

        all_scores.sort(key=operator.itemgetter(0))
        sorted_scores = sorted(all_scores, reverse=True, key=operator.itemgetter(1))
        rows.append(exp_vars + (len(all_scores),) + (sorted_scores[0] if sorted_scores else (-1, -1)))
        accuracy_per_run[exp_name] = {
            'x': [s[0] for s in all_scores],
            'all': [s[1] for s in all_scores],
        }
        print(d)

    metric_types = tuple(sorted(metric_types))
    df = pd.DataFrame(rows, columns=('att', 'num steps eval', 'step', 'exact match'))
    flat_df = pd.DataFrame(
        [(exp_name, step) + tuple(metrics.get(t) for t in metric_types) for exp_name, step, metrics in all_metrics],
        columns=('exp_name', 'step') + metric_types)
    
    return flat_df


baseline_df = find_baseline_checkpoints()
print(baseline_df.loc[baseline_df['corpus BLEU'].idxmax()])
print(baseline_df.loc[baseline_df['exact match'].idxmax()])

# # Loading ground truth data

# +
import ast
import astor 

val_data = registry.construct(
    'dataset', json.loads(_jsonnet.evaluate_file('configs/hearthstone/nl2code.jsonnet'))['data']['val']) 
test_data = registry.construct(
    'dataset', json.loads(_jsonnet.evaluate_file('configs/hearthstone/nl2code.jsonnet'))['data']['test']) 

normalized_gold_code_val = [astor.to_source(ast.parse(item.code)) for item in val_data]
normalized_gold_code_test = [astor.to_source(ast.parse(item.code)) for item in test_data]
# -

# # Determining when and where idioms were output during decoding

# +
idiom_exact_match_max = idiom_df.loc[idiom_df['exact match'].idxmax()]

idiom_exact_match_infer_val_history_path = \
    'logdirs/20190201-hs-allmatches-anysplit-multimean/{}/infer-{}-step{:05d}-bs1-with-history.jsonl'.format(
        idiom_exact_match_max['exp_name'], 'val', idiom_exact_match_max['step'])

# !CUDA_VISIBLE_DEVICES= python infer.py \
#     --config configs/hearthstone-idioms/nl2code-0201-allmatches-anysplit-multimean.jsonnet \
#     --logdir logdirs/20190201-hs-allmatches-anysplit-multimean/filt-contains-hole_st-cov-examples_nt-80 \
#     --config-args "{filt: 'contains-hole', st: 'cov-examples', nt: 80}" \
#     --output-history \
#     --output __LOGDIR__/infer-val-step02200-bs1-with-history.jsonl --step 2200 --section val --beam-size 1

idiom_exact_match_infer_test_history_path = \
    'logdirs/20190201-hs-allmatches-anysplit-multimean/{}/infer-{}-step{:05d}-bs1-with-history.jsonl'.format(
        idiom_exact_match_max['exp_name'], 'test', idiom_exact_match_max['step'])

# !CUDA_VISIBLE_DEVICES= python infer.py \
#     --config configs/hearthstone-idioms/nl2code-0201-allmatches-anysplit-multimean.jsonnet \
#     --logdir logdirs/20190201-hs-allmatches-anysplit-multimean/filt-contains-hole_st-cov-examples_nt-80 \
#     --config-args "{filt: 'contains-hole', st: 'cov-examples', nt: 80}" \
#     --output-history \
#     --output __LOGDIR__/infer-test-step02200-bs1-with-history.jsonl --step 2200 --section test --beam-size 1

# +
all_rules = json.load(
    open('data/hearthstone-idioms-20190201/all-matches-trees-anysplit/filt-contains-hole_st-cov-examples_nt-80/nl2code/grammar_rules.json')
)['all_rules']

def count_nodes(tree):
    queue = collections.deque([tree])
    count = 0
    while queue:
        node = queue.pop()
        count += 1 
        if isinstance(node, dict):
            for k, v in node.items():
                if k == '_type':
                    continue
                if isinstance(v, (list, tuple)):
                    queue.extend(v)
                else:
                    queue.append(v)
    return count

def analyze_idiom_usage(infer_history_path, normalized_gold_code):
    inferred_all = [json.loads(line) for line in open(infer_history_path)]
    exact_match_all = [
        normalized_gold_code[i] == example['beams'][0]['inferred_code'] if example['beams'] else False
        for i, example in enumerate(inferred_all)
    ]
    
    decoding_history_all = {}
    for example in inferred_all:
        decoding_history = []
        decoding_history_all[example['index']] = decoding_history
        if not example['beams']:
            continue

        for choice in example['beams'][0]['choice_history']:
            if isinstance(choice, int):
                decoding_history.append(all_rules[choice])
            else:
                decoding_history.append(choice)
                
    # Number of idioms used per program
    num_idioms_used_all = {}
    for i, history in decoding_history_all.items():
        counter = collections.Counter()
        for choice in history:
            if isinstance(choice, list) and isinstance(choice[1], str) and re.match('Template\d+', choice[1]):
                counter[tuple(choice)] += 1
        num_idioms_used_all[i] = counter
    
    # Size of each program
    # 1. Number of characters
    # 2. Number of lines
    # 3. Number of nodes in tree
    num_characters_all = {example['index']: len(example['beams'][0]['inferred_code']) if example['beams'] else 0 for example in inferred_all}
    num_lines_all      = {example['index']: example['beams'][0]['inferred_code'].count('\n') if example['beams'] else 0 for example in inferred_all}
    num_nodes_all      = {example['index']: count_nodes(example['beams'][0]['model_output']) if example['beams'] else 0 for example in inferred_all}
    
    # Dataframe
    idiom_usage_df = pd.DataFrame(collections.OrderedDict((
        ('Number of idioms used', {k: sum(v.values()) for k, v in num_idioms_used_all.items()}),
        ('Number of characters', num_characters_all),
        ('Number of lines', num_lines_all),
        ('Number of AST nodes', num_nodes_all),
        ('Exact match', dict(enumerate(exact_match_all))),
    )))
    exact_match_idiom_usage_df = idiom_usage_df[idiom_usage_df['Exact match']]
    
    fig, ax = plt.subplots()
    sns.distplot(idiom_usage_df['Number of idioms used'], kde=False, ax=ax, label='All')
    
    fig, ax = plt.subplots()
    sns.distplot(exact_match_idiom_usage_df['Number of idioms used'], kde=False, ax=ax, bins=5, label='Exact match')
    

    for x, y in (
        ('Number of characters', 'Number of idioms used'),
        ('Number of AST nodes', 'Number of idioms used')
    ):
#        fig, ax = plt.subplts
#        plt.plot(x, y, data=idiom_usage_df)
        jg = sns.jointplot(x=x, y=y, data=idiom_usage_df)
        jg.fig.suptitle('All')
        
        jg = sns.jointplot(x=x, y=y, data=exact_match_idiom_usage_df)
        jg.fig.suptitle('Exactly matched')
    
        
#        sns.jointplot(x='Number of characters', y='Number of idioms used', data=idiom_usage_df),
#        sns.jointplot(x='Number of characters', y='Number of idioms used', data=exact_match_idiom_usage_df),
#        sns.jointplot(x='Number of lines', y='Number of idioms used', data=idiom_usage_df),
#        sns.jointplot(x='Number of lines', y='Number of idioms used', data=exact_match_idiom_usage_df),
#        sns.jointplot(x='Number of AST nodes', y='Number of idioms used', data=idiom_usage_df),
#        sns.jointplot(x='Number of AST nodes', y='Number of idioms used', data=exact_match_idiom_usage_df),
#    ]:
        
# -

analyze_idiom_usage(idiom_exact_match_infer_val_history_path, normalized_gold_code_val)

analyze_idiom_usage(idiom_exact_match_infer_test_history_path, normalized_gold_code_test)


# # Teacher-forced precision and recall

# +
def analyze_teacher_forced_pr(report, templates, accuracy_ks=(1, 2), precision_ks=(1, 2), recall_ks=(1, 2)):
    # Count how often 
    template_match_counts = collections.defaultdict(int)
    template_choice_ranks = {
        'all': collections.defaultdict(list),
        'templates only': collections.defaultdict(list),        
    }
    template_valid_choice_ranks = {
        'all': collections.defaultdict(list),
        'templates only': collections.defaultdict(list),
    }
    min_valid_ranks = []
    
    for item in report:
        for entry in item['history']:
            if not isinstance(entry['choices'][0], str):
                continue
            
            all_ranks = {}
            template_only_ranks = {}
            
            template_only_i = 0
            for i, choice in enumerate(entry['choices']):
                all_ranks[choice] = i
                template_only_ranks[choice] = template_only_i
                if not re.match('Template(\d+).*', choice):
                    template_only_i += 1
            
            # For overall top-k accuracy
            min_valid_rank = min(all_ranks[choice] for choice in entry['valid_choices'])
            min_valid_ranks.append(min_valid_rank)
            
            # For precision
            #   times that choice appeared at rank k, when it was valid (template_valid_choice_ranks)
            # / times that choice appeared at rank k (tp + fp), whether or not valid (template_choice_ranks)
            for choice in entry['choices']:
                m = re.match('Template(\d+).*', choice)
                if not m:
                    continue
                template_id = int(m.group(1))
                template_choice_ranks['all'][template_id].append(all_ranks[choice])
                template_choice_ranks['templates only'][template_id].append(template_only_ranks[choice])
                
            # For recall
            #   times that choice appeared at rank k, when it was valid (template_valid_choice_ranks)
            # / times that choice was valid (tp + fn)
            for choice in entry['valid_choices']:                
                m = re.match('Template(\d+).*', choice)
                if not m:
                    continue
                template_id = int(m.group(1))
                template_match_counts[template_id] += 1
                
                # Determine its rank
                template_valid_choice_ranks['all'][template_id].append(all_ranks[choice])
                # Determine its rank, excluding other templates
                template_valid_choice_ranks['templates only'][template_id].append(template_only_ranks[choice])

    min_valid_ranks = np.array(min_valid_ranks)
    # Top-k accuracy: there exists a valid choice such that its rank ≤ k
    top_k_accuracy = {
        k: np.sum(min_valid_ranks < k) / len(min_valid_ranks) for k in accuracy_ks
    }
    
    # Precision:
    #   Among the # of times that a given template has rank ≤ k, 
    #   how often it is a valid choice
    top_k_precision = {
        type_name: {
            k: {i: np.sum(np.array(template_valid_choice_ranks[type_name][i]) < k) /
                   np.sum(np.array(template_choice_ranks[type_name][i]) < k)
                for i in template_match_counts.keys()} for k in precision_ks
        } for type_name in template_valid_choice_ranks
    }
    
    # Recall:
    #   Among the # of times that a given template is a valid choice,
    #   how often it is a choice with rank ≤ k
    top_k_recall = {
        type_name: {
            k: {i: np.sum(np.array(ranks) < k) / len(ranks) for i, ranks in ranks_of_type.items()} for k in recall_ks
        } for type_name, ranks_of_type in template_valid_choice_ranks.items()}

    accuracy_df = pd.DataFrame({
        'Accuracy @ {}'.format(k): [top_k_accuracy[k]]
        for k in accuracy_ks
    })
    pr_df = pd.DataFrame({
        'Head': {t['id']: t['idiom'][0] for t in templates},
        'Matches': template_match_counts, 
        **{
            'Precision @ {} {}'.format(k, type_name): top_k_precision[type_name][k]
            for type_name in top_k_precision.keys()
            for k in precision_ks
        },
        **{
            'Recall @ {} {}'.format(k, type_name): top_k_recall[type_name][k]
            for type_name in top_k_recall.keys()
            for k in recall_ks
        }
    })
    return accuracy_df, pr_df

def analyze_anysplit_one(name, section):
    report = [json.loads(line)
              for line
              in open('../logdirs/20190201-hs-allmatches-anysplit/{}/debug-{}-step2600.jsonl'.format(name, section))]
    templates = json.load(open('../data/hearthstone-idioms-20190201/all-matches-trees-anysplit/{}/templates.json'.format(name)))
    return analyze(report, templates)


# +
idiom_exact_match_debug_val_path = \
    'logdirs/20190201-hs-allmatches-anysplit-multimean/{}/debug-{}-step{:05d}.jsonl'.format(
        idiom_exact_match_max['exp_name'], 'val', idiom_exact_match_max['step'])

# !CUDA_VISIBLE_DEVICES= python infer.py \
#     --mode debug \
#     --config configs/hearthstone-idioms/nl2code-0201-allmatches-anysplit-multimean.jsonnet \
#     --logdir logdirs/20190201-hs-allmatches-anysplit-multimean/filt-contains-hole_st-cov-examples_nt-80 \
#     --config-args "{filt: 'contains-hole', st: 'cov-examples', nt: 80}" \
#     --output __LOGDIR__/debug-val-step02200.jsonl --step 2200 --section val --beam-size 1

idiom_exact_match_debug_test_path = \
    'logdirs/20190201-hs-allmatches-anysplit-multimean/{}/debug-{}-step{:05d}.jsonl'.format(
        idiom_exact_match_max['exp_name'], 'test', idiom_exact_match_max['step'])

# !CUDA_VISIBLE_DEVICES= python infer.py \
#     --mode debug \
#     --config configs/hearthstone-idioms/nl2code-0201-allmatches-anysplit-multimean.jsonnet \
#     --logdir logdirs/20190201-hs-allmatches-anysplit-multimean/filt-contains-hole_st-cov-examples_nt-80 \
#     --config-args "{filt: 'contains-hole', st: 'cov-examples', nt: 80}" \
#     --output-history \
#     --output __LOGDIR__/debug-test-step02200.jsonl --step 2200 --section test --beam-size 1
# -

templates = json.load(
    open('data/hearthstone-idioms-20190201/all-matches-trees-anysplit/{}/templates.json'.format(idiom_exact_match_max['exp_name'])))
accuracy_df_val, pr_df_val = analyze_teacher_forced_pr([json.loads(line) for line in open(idiom_exact_match_debug_val_path)], templates)

accuracy_df_val

pr_df_val

accuracy_df_test, pr_df_test = analyze_teacher_forced_pr([json.loads(line) for line in open(idiom_exact_match_debug_test_path)], templates)

accuracy_df_test

pr_df_test


# # Comparing versus baseline

def load_inferred(infer_history_path, normalized_gold_code):
    inferred_all = [json.loads(line) for line in open(infer_history_path)]
    exact_match_all = [
        normalized_gold_code[i] == example['beams'][0]['inferred_code'] if example['beams'] else False
        for i, example in enumerate(inferred_all)
    ]
    
    return inferred_all, exact_match_all


print(baseline_df.loc[baseline_df['exact match'].idxmax()])

# +
# !CUDA_VISIBLE_DEVICES= python infer.py \
#     --config configs/hearthstone/nl2code.jsonnet \
#     --logdir logdirs/20181231-nl2code-hearthstone-fef2c5b/att1 \
#     --output __LOGDIR__/infer-test-step02300-bs1.jsonl --step 2300 --section test --beam-size 1

baseline_exact_match_infer_test_path = 'logdirs/20181231-nl2code-hearthstone-fef2c5b/att1/infer-test-step02300-bs1.jsonl'
# -

inferred_test, exact_match_test = load_inferred(idiom_exact_match_infer_test_history_path, normalized_gold_code_test)
baseline_inferred_test, baseline_exact_match_test = load_inferred(baseline_exact_match_infer_test_path, normalized_gold_code_test)

comparison_df = pd.DataFrame({'Idioms': exact_match_test, 'Baseline': baseline_exact_match_test})
comparison_df[comparison_df.Baseline != comparison_df.Idioms]

from seq2struct import grammars
idiom_config = json.loads(_jsonnet.evaluate_file(
        'configs/hearthstone-idioms/nl2code-0201-allmatches-anysplit-multimean.jsonnet',
        tla_codes={'args': "{filt: 'contains-hole', st: 'cov-examples', nt: 80}"}))
idiom_grammar = registry.construct('grammar', idiom_config['model']['decoder_preproc']['grammar'])

import pprint
for i in comparison_df[comparison_df.Baseline != comparison_df.Idioms].index:
    print('#{}\nGold:\n{}\nIdioms:\n{}\nBaseline:\n{}'.format(
        i, normalized_gold_code_test[i], 
        inferred_test[i]['beams'][0]['inferred_code'],
        baseline_inferred_test[i]['beams'][0]['inferred_code']))
    
    templates_used = set()
    actions = []
    for choice in inferred_test[i]['beams'][0]['choice_history']:
        if isinstance(choice, int):
            rule = all_rules[choice]
            actions.append(rule)
            if isinstance(rule[1], str) and re.match('Template\d+', rule[1]):
                templates_used.add(rule[1])
        else:
            actions.append(choice)
    if i == 36:
        print('Idiom actions:\n{}\n'.format(pprint.pformat(actions)))
    
    print('Templates used:')
    for t in natsort.natsorted(templates_used):
        cons = idiom_grammar.ast_wrapper.constructors[t]
        tree = cons.template({
            i: ['<HOLE-{}>*'.format(i)] if field.seq else '<HOLE-{}>'.format(i)
            for i, field in enumerate(cons.fields)}
        )
        print(t)
        pprint.pprint(tree)
    print('---')



StarcoderdataJupyterScriptsDedupFiltered