# Copyright (c) 2014-2015, NVIDIA CORPORATION.  All rights reserved.

import os
import io
import re
import json
import math
import tarfile
import zipfile
from collections import OrderedDict
from datetime import timedelta

import flask
import werkzeug.exceptions


import digits
from digits import utils
from digits.webapp import app, scheduler, autodoc
from digits.utils import time_filters
from digits.utils.routing import request_wants_json
from . import ModelJob
import forms
import images.views
import images as model_images

from digits import frameworks

NAMESPACE = '/models/'

@app.route(NAMESPACE, methods=['GET'])
@autodoc(['models'])
def models_index():
    column_attrs = list(get_column_attrs())
    raw_jobs = [j for j in scheduler.jobs if isinstance(j, ModelJob)]

    column_types = [
        ColumnType('latest', False, lambda outs: outs[-1]),
        ColumnType('max', True, lambda outs: max(outs)),
        ColumnType('min', True, lambda outs: min(outs))
    ]

    jobs = []
    for rjob in raw_jobs:
        train_outs = rjob.train_task().train_outputs
        val_outs = rjob.train_task().val_outputs
        history = rjob.status_history

        # update column attribute set
        keys = set(train_outs.keys() + val_outs.keys())

        # build job dict
        job_info = JobBasicInfo(
            rjob.name(),
            rjob.id(),
            rjob.status,
            time_filters.print_time_diff_nosuffixes(history[-1][1] - history[0][1]),
            rjob.train_task().framework_id
        )

        # build a dictionary of each attribute of a job. If an attribute is
        # present, add all different column types.
        job_attrs = {}
        for cattr in column_attrs:
            if cattr in train_outs:
                out_list = train_outs[cattr].data
            elif cattr in val_outs:
                out_list = val_outs[cattr].data
            else:
                continue

            job_attrs[cattr] = {ctype.name: ctype.find_from_list(out_list)
                for ctype in column_types}

        job = (job_info, job_attrs)
        jobs.append(job)

    attrs_and_labels = []
    for cattr in column_attrs:
        for ctype in column_types:
            attrs_and_labels.append((cattr, ctype, ctype.label(cattr)))

    return flask.render_template('models/index.html',
        jobs=jobs,
        attrs_and_labels=attrs_and_labels)

@app.route(NAMESPACE + '<job_id>.json', methods=['GET'])
@app.route(NAMESPACE + '<job_id>', methods=['GET'])
@autodoc(['models', 'api'])
def models_show(job_id):
    """
    Show a ModelJob

    Returns JSON when requested:
        {id, name, directory, status, snapshots: [epoch,epoch,...]}
    """
    job = scheduler.get_job(job_id)
    if job is None:
        raise werkzeug.exceptions.NotFound('Job not found')

    if request_wants_json():
        return flask.jsonify(job.json_dict(True))