import os
import csv
import pandas as pd
import numpy as np

# A magic number the _could be read/inferrred from the data, but let's
# be real: It's not changing.
NUM_CUB_CLASSES = 200
NUM_CUB_ATTRIBUTES = 312

def load_attribute_schema(dataset_path='',clean=True):
# Return dict with key as the attribute id (int type) and value a
# 2-tuple (attribute,value) where attribute is the description of the
# attribute and value the specific value it takes, ie.,
# ('has_bill_shape','dagger') clean indicates whether to alter the
# values for use in textual descriptions according to pre-codified
# rules.
    
    attributes = {}
    with open(os.path.join(dataset_path,'attributes.txt')) as fd:
        # Each line has the form:
        # 12 has_wing_color::iridescent
        for attr_line in fd:
            [attr_id,attr_name_val] = attr_line.strip().split(' ')
            attributes[int(attr_id)] = attr_name_val.split('::')

    if clean: # Hacky list of alterations to make for captioning
        # has_bill_shape
        attributes[1][1] = 'curved' # curved_(up_or_down)

        # has_tail_shape
        for a in range(74,80): # 74:79
            attributes[a][1] = attributes[a][1].removesuffix('_tail')

        # has_bill_length
        attributes[150][1] = '' # about_the_same_as_head
        attributes[151][1] = 'long' # longer_than_head
        attributes[152][1] = 'short' # shorter_than_head

        # has_wing_shape
        for a in range(213,218): # 213:217
            attributes[a][1] = attributes[a][1].removesuffix('-wings')
            
    return attributes
            
def load_class_attributes(dataset_path=''):
# Return Pandas DataFrame with indices being the CUB class id and the
# columns being attribute IDs. Ebtry value is the percentage of the
# time (0-100) the class is imputed with that attribute.

    attr_arr = np.loadtxt(
        os.path.join(
            dataset_path,'attributes/class_attribute_labels_continuous.txt'))

    return pd.DataFrame(attr_arr,
                        index=list(range(1,NUM_CUB_CLASSES+1)),
                        columns=list(range(1,NUM_CUB_ATTRIBUTES+1)) )

    
                 
class CUB():

    def __init__(self,dataset_path='', map_file=None, attr_thresh=50.0):

        self.attr_thresh = attr_thresh
        
        # Load in the attributes schema
        self.attr_schema = load_attribute_schema(dataset_path)
        
        # Load the attribute matrix
        self.attributes = load_class_attributes(dataset_path)

        # Invert to access e.g., 'has_eye_color' --> [136, 137, ...]
        self.invert_attribute_index()

        # Load (species) conversion from NABird ID to CUB ID
        if map_file is not None:
            with open(map_file) as fd:
                self.nab_map = { int(row['nabird_id']):int(row['cub_id'])
                                 for row in csv.DictReader(fd) }


    def get_id(self,nabird_id):
        # Produce the CUB id (int)
        if nabird_id in self.nab_map:
            return self.nab_map[nabird_id]
        else:
            return None
    
    def invert_attribute_index(self):
        # Produce/store a dictionary mapping attribute name to a list
        # of integers (cub_attr_id values) for that name. I.e.,
        # {'has_bill_shape': [1,2,3,4,5,6,7,8,9], ... }

        inv_index = {}

        # Iterate over all attribute ids, re-indexing by names (building lists)
        for (attr_id,attr_name_val) in self.attr_schema.items():

            [attr_name,attr_val] = attr_name_val
            
            if attr_name in inv_index: # key already added?
                inv_index[attr_name].append(attr_id) # append id to list
            else:
                inv_index[attr_name] = [attr_id] # initialize list with id

        # Ensure a canonical ordering by sorting names in list values
        for attr_names in inv_index.values():
            attr_names.sort() 
        
        self.attr_index = inv_index
                
        
    def attribute_values(self,bird_id,attr_name):
        # Get the list of an attribute's descriptions for a given bird species

        attr_ids = self.attr_index[attr_name] # list of indices for attribute

        # Pandas magic
        bird_attr = self.attributes.loc[bird_id]  # row for given bird
        #print('row',bird_attr)
        bird_attr = bird_attr[attr_ids]           # cols for given attribute name
        #print('cols',bird_attr)
        bird_attr = bird_attr > self.attr_thresh  # find specific attr values
        #print('matches',bird_attr)
        bird_attr = bird_attr.index[bird_attr]    # extract matching values
        #print('indices',bird_attr)

        # extract names of matching values
        names = [self.attr_schema[attr_id][1] for attr_id in bird_attr]

        return names
