import os

import imageio
import numpy as np
import pandas as pd

newt_labels_fp = "newt2021_labels.csv"
newt_image_dir = "newt2021_images/"

newt_df = pd.read_csv(newt_labels_fp)
newt_df['filepath'] = newt_df['id'].apply(
    lambda image_id: os.path.join(newt_image_dir, image_id + ".jpg")
)

task_results = []
for task_name, task_df in newt_df.groupby('task'):
    
    X_train = []
    X_test = []
    y_train = []
    y_test = []
    for _, image_row in task_df.iterrows():
        
        image = imageio.imread(image_row['filepath'])
        image_label = image_row['label']

        # Place holder for feature extraction
        image_feature = np.random.default_rng().uniform(-1,1,1000) 
        
        if image_row['split'] == 'train':
            X_train.append(image_feature)
            y_train.append(image_label)
        else:
            X_test.append(image_feature)
            y_test.append(image_label)
    
    X_train = np.array(X_train)
    X_test = np.array(X_test)
    y_train = np.array(y_train)
    y_test = np.array(y_test)
    
    # Placeholder for training and evaluating a classifier
    task_acc = np.random.default_rng().uniform(0,1,1)[0]
    
    task_results.append({
        'task' : task_name,
        'task_cluster' : task_df.iloc[0]['task_cluster'],
        'task_subcluster' : task_df.iloc[0]['task_subcluster'],
        'acc' : task_acc
    })

task_results_df = pd.DataFrame(task_results)
print(f"Mean Accuracy: {task_results_df['acc'].mean():.3f}")

for task_cluster, mean_acc in task_results_df.groupby('task_cluster')['acc'].mean().iteritems():
    print(f"{task_cluster:10} {mean_acc:0.3f}")

import os

import pandas as pd

####
# Load in the NeWT dataset
newt_labels_fp = "newt2021_labels.csv"
newt_image_dir = "newt2021_images/"

newt_df = pd.read_csv(newt_labels_fp)
newt_df['filepath'] = newt_df['id'].apply(
    lambda image_id: os.path.join(newt_image_dir, image_id + ".jpg")
)

####
# Print some basic info for the task clusters
newt_task_clusters = newt_df['task_cluster'].unique()
print("NeWT Task Clusters")
for i, task_cluster_name in enumerate(newt_task_clusters):
    task_cluster_df = newt_df[newt_df['task_cluster'] == task_cluster_name]
    task_subclusters = task_cluster_df['task_subcluster'].dropna().unique()
    num_subclusters = task_subclusters.shape[0]
    num_tasks = task_cluster_df['task'].unique().shape[0]
    print(f"{i+1:1d}. {task_cluster_name:10s} ({num_subclusters:2d} subclusters) : {num_tasks:3d} tasks")
    if num_subclusters > 0:
        for task_subcluster_name in task_subclusters:
            task_subcluster_df = task_cluster_df[task_cluster_df['task_subcluster'] == task_subcluster_name]
            num_tasks = task_subcluster_df['task'].unique().shape[0]
            print(f"\t{task_subcluster_name:10s}: {num_tasks:3d} tasks")
    print()
print()
        
####
# Choose a specific task cluster and print basic info about the tasks
task_cluster_name = 'appearance'
assert task_cluster_name in newt_task_clusters, f"{task_cluster_name} is not a valid NeWT task cluster"
task_cluster_df = newt_df[newt_df['task_cluster'] == task_cluster_name]

num_subclusters = task_cluster_df['task_subcluster'].dropna().unique().shape[0]
num_tasks = task_cluster_df['task'].unique().shape[0]

if num_subclusters > 0:
    print(f"The task `{task_cluster_name}` has {num_tasks:3d} total tasks, grouped into {num_subclusters} subclusters")
else:
    print(f"The task `{task_cluster_name}` has {num_tasks:3d} total tasks")
    

print(f"Tasks included in the `{task_cluster_name}` cluster:")
for i, (task_name, task_df) in enumerate(task_cluster_df.groupby('task')):
    
    print(f"{i+1:3d}. {task_name:55s}")
    
    # Get positive and negative examples for this task
    task_positive_df = task_df[task_df['label'] == 1]
    task_negative_df = task_df[task_df['label'] == 0]
    
    # Get the text label for the positive and negative classes
    # NOTE: these are not unique across the tasks. 
    positive_text_label = task_positive_df.iloc[0]['text_label']
    negative_text_label = task_negative_df.iloc[0]['text_label']
    print(f"\tpositive class text label: {positive_text_label}")
    print(f"\tnegative class text label: {negative_text_label}")
    print(f"\tpositive images: {task_positive_df.shape[0]:3d}, negative images: {task_negative_df.shape[0]:3d}")
    print()
    
    # Get `train` and `test` examples for this task
    for split in ['train', 'test']:
        task_split_df = task_df[task_df['split'] == split]
        
        task_split_pos_df = task_split_df[task_split_df['label'] == 1]
        task_split_neg_df = task_split_df[task_split_df['label'] == 0]
        
        print(f"\tpositive {split:5s} images: {task_split_pos_df.shape[0]:3d}")
        print(f"\tnegative {split:5s} images: {task_split_neg_df.shape[0]:3d}")
        print()