# Script to load all scraped levels and compute features.

import csv, json, os, sys
import datetime
import numpy as np
from queue import Queue
from termcolor import cprint

import constants
from features.segment_graph import SegmentGraph
from features.state_graph import Graph

fieldnames = constants.FIELDNAMES

# Recompute only given features.
comp_features = set()
if len(sys.argv) > 1:
    comp_features = set()
    for feature in sys.argv[1:]:
        if feature not in fieldnames:
            cprint('Feature "%s" does not exist.' % feature, 'red', attrs=['bold'])
            quit()
        comp_features.add(feature)
    cprint('Recomputing features: %s' % ', '.join(comp_features), 'green', attrs=['bold'])
else:
    cprint('Recomputing all features: %s' % ', '.join(fieldnames), 'green', attrs=['bold'])

# Get levels and logs directories.
current_dir = os.getcwd()
data_dir = 'data'.join(current_dir.rsplit('code', 1)) + '/'
levels_dir = os.path.join(data_dir, 'levels')
all_logs_dir = os.path.join(current_dir, 'logs')
log_dir = os.path.join(all_logs_dir, datetime.datetime.now().strftime('%Y%m%d'))
log_name = 'features_%s.out' % datetime.datetime.now().strftime('%H%M%S')

if not os.path.exists(log_dir):
    os.makedirs(log_dir)

feature_names, X, X_ids = None, None, None
if len(comp_features) > 0:
    feature_cnt = constants.FEATURE_CNT
    feature_file = os.path.join(data_dir, constants.FEATURES_FILE)
    X = np.genfromtxt(feature_file, delimiter=',', dtype='|S255')
    feature_names = [x.decode('UTF-8') for x in X[0,0:feature_cnt+1]]
    X = np.genfromtxt(feature_file, delimiter=',', skip_header=1)
    X_ids = [int(x) for x in X[:, 0]]

class Level:
    '''Class representing level.'''
    def __init__(self, level, id):
        isCheckpoint = level['isCheckpoint']
        hasWall = [
            level['hasTopWall'],
            level['hasRightWall'],
            level['hasBottomWall'],
            level['hasLeftWall']
        ]
        ball_row = level['ball_row']
        ball_col = level['ball_col']

        self.checks_count = sum((map(sum, isCheckpoint)))
        self.width = level['width']
        self.height = level['height']
        self.id = id

        # Creating both graphs.
        self.segment_graph = SegmentGraph(hasWall, isCheckpoint, ball_row, ball_col)
        self.graph = Graph(hasWall, isCheckpoint, ball_row, ball_col)

        self.feature_map = {
            'id': lambda: self.id,
            'checkpoint_count': lambda: self.checks_count,
            'tile_count': lambda: self.width * self.height,
            'width': lambda: self.width,
            'reachable_tile_count': self.graph.get_reachable_tile_count,
            'scc_count': self.segment_graph.get_scc_count,
            'scc_checkpoint_count': self.segment_graph.get_scc_checkpoint_count,
            'shortest_path': self.graph.get_shortest_path,
            'reachable_states_count': self.graph.get_reachable_states_count,
            'viable_states_count': self.graph.get_viable_states_count,
            'shortest_path_tiles': self.graph.get_shortest_path_tiles,
        }
        self.type_map = {
            'id': int,
            'checkpoint_count': int,
            'tile_count': int,
            'width': int,
            'reachable_tile_count': int,
            'scc_count': int,
            'scc_checkpoint_count': int,
            'shortest_path': int,
            'reachable_states_count': int,
            'viable_states_count': int,
            'shortest_path_tiles': int,
        }

    def get_feature(self, f_name):
        if f_name not in self.feature_map:
            return None
        if len(comp_features) > 0 and f_name not in comp_features:
            return self.type_map[f_name](X[X_ids.index(self.id), feature_names.index(f_name)])
        result = self.feature_map[f_name]()
        if result is None:
            return None
        return self.type_map[f_name](result)

def load_level(level_name):
    '''Load level to dictionary.'''
    with open(level_name, 'r') as input:
        level_text = input.read()
        return json.loads(level_text)

if __name__ == '__main__':

    Y = np.genfromtxt(os.path.join(data_dir, constants.USER_TIME_FILE), delimiter=',', dtype=int)
    # First column contains user names.
    Y = np.delete(Y, 0, 1)
    # First line contains names of levels.
    levels_ids = Y[0]
    Y = Y[1:]

    # Compute avg, avg_log, 25per (from log) and std of all columns.
    mask_Y  = np.ma.array(Y, mask=Y < 0)
    av = np.ma.average(mask_Y, axis=0)
    st = np.ma.std(mask_Y, axis = 0)
    cprint('Ignore next warning:', 'yellow', attrs=['bold'])
    log_Y = np.log(Y)
    masked_log_Y = np.ma.masked_invalid(log_Y)
    av_log = np.ma.average(masked_log_Y, axis=0)
    per = np.nanpercentile(log_Y, 25, axis=0)

    # For every level, compute features, avg time, avg_log, 25per time and std.
    for i in range(len(levels_ids)):
        level_chars = {}
        level_id = levels_ids[i]
        cprint('%3d: Level %s : %s' % (i, level_id, datetime.datetime.now()), 'green')
        # Just save those, which are already computed.
        level_chars['avg_log_time'] = av_log[i]
        level_chars['avg_time'] = av[i]
        level_chars['stdev'] = st[i]
        level_chars['25per'] = per[i]

        level_name = os.path.join(levels_dir, constants.get_level_name(level_id))
        level = Level(load_level(level_name), level_id)
        # Ask Level for everything else.
        for f_name in fieldnames:
            res = level.get_feature(f_name)
            if res is None: continue
            level_chars[f_name] = res
        with open(os.path.join(log_dir, log_name), 'a+') as out:
            out.write('%s\n' % json.dumps(level_chars))

    chars = []
    with open(os.path.join(log_dir, log_name), 'r') as input:
        for line in input:
            level_text = line.strip()
            chars.append(json.loads(level_text))

    # Magicly write output to csv.
    with open(os.path.join(data_dir, constants.FEATURES_FILE), 'w', newline='') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames, dialect='unix', quoting=csv.QUOTE_NONE)

        writer.writeheader()
        for level in chars:
            writer.writerow(level)

