# Authors: Robert Luke <mail@robertluke.net>
#
# License: BSD (3-clause)

import pandas as pd
import statsmodels.formula.api as smf

from ...simulation import simulate_nirs_raw
from ...experimental_design import make_first_level_design_matrix
from ...statistics import run_GLM, statsmodels_to_results
from ...utils._io import glm_to_tidy

def test_statsmodel_to_df():

    amplitude = 1.432

    df_cha = pd.DataFrame()
    for n in range(5):

        raw = simulate_nirs_raw(sfreq=3., amplitude=amplitude,
                                sig_dur=300., stim_dur=5.,
                                isi_min=15., isi_max=45.)
        design_matrix = make_first_level_design_matrix(raw, stim_dur=5.0)
        glm_est = run_GLM(raw, design_matrix)
        cha = glm_to_tidy(raw, glm_est, design_matrix)
        cha["ID"] = '%02d' % n
        df_cha = df_cha.append(cha)
    df_cha["theta"] = df_cha["theta"] * 1.0e6
    roi_model = smf.mixedlm("theta ~ -1 + Condition", df_cha,
                            groups=df_cha["ID"]).fit(method='nm')
    df = statsmodels_to_results(roi_model)
    assert type(df) == pd.DataFrame
    assert df["Coef."]["Condition[A]"] == amplitude
    assert df["Significant"]["Condition[A]"]
    assert df.shape == (8, 8)

    roi_model = smf.rlm("theta ~ -1 + Condition", df_cha,
                        groups=df_cha["ID"]).fit()
    df = statsmodels_to_results(roi_model)
    assert type(df) == pd.DataFrame
    assert df["Coef."]["Condition[A]"] == amplitude
    assert df["Significant"]["Condition[A]"]
    assert df.shape == (8, 8)


import os

import numpy as np
import torch

from instanceseg.datasets import dataset_generator_registry, sampler, dataset_registry
from instanceseg.utils.misc import pop_without_del
from instanceseg.factory.samplers import get_configured_sampler

DEBUG_ASSERTS = True

def get_datasets_with_transformations(dataset_type, cfg, transform=True):
    train_dataset, val_dataset = dataset_generator_registry.get_dataset(dataset_type, cfg, transform)
    return train_dataset, val_dataset

def get_dataloaders(cfg, dataset_type, cuda, sampler_cfg=None):
    # 1. dataset
    train_dataset, val_dataset = get_datasets_with_transformations(dataset_type, cfg, transform=True)

    # 2. samplers
    if sampler_cfg is not None and isinstance(sampler_cfg['val'], str) and sampler_cfg['val'] == 'copy_train':
        val_dataset = train_dataset
    train_sampler, val_sampler, train_for_val_sampler = get_samplers(dataset_type, sampler_cfg,
                                                                     train_dataset, val_dataset)

    # Create dataloaders from datasets and samplers
    loader_kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, sampler=train_sampler, **loader_kwargs)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, sampler=val_sampler, **loader_kwargs)
    train_loader_for_val = torch.utils.data.DataLoader(train_dataset, batch_size=1,
                                                       sampler=train_for_val_sampler, **loader_kwargs)

    if DEBUG_ASSERTS:
        try:
            i, [sl, il] = [d for i, d in enumerate(train_loader) if i == 0][0]
        except:
            raise

    return {
        'train': train_loader,
        'val': val_loader,
        'train_for_val': train_loader_for_val,
    }

def get_samplers(dataset_type, sampler_cfg, train_dataset, val_dataset):

    if sampler_cfg is None:
        train_sampler = sampler.sampler.RandomSampler(train_dataset)
        val_sampler = sampler.sampler.SequentialSampler(val_dataset)
        train_for_val_sampler = sampler.sampler.SequentialSampler(train_dataset)
    else:
        # Get 'clean' datasets for instance counting
        default_train_dataset, default_val_dataset, transformer_tag = \
            dataset_generator_registry.get_default_datasets_for_instance_counts(dataset_type)

        # train sampler
        train_sampler = sampler_generator_helper(dataset_type, train_dataset, default_train_dataset,
                                                 sampler_cfg, 'train', transformer_tag)

        # val sampler
        if isinstance(sampler_cfg['val'], str) and sampler_cfg['val'] == 'copy_train':
            val_sampler = train_sampler.copy(sequential_override=True)
        else:
            val_sampler = sampler_generator_helper(dataset_type, val_dataset, default_val_dataset,
                                                   sampler_cfg, 'val', transformer_tag)

        # train_for_val sampler
        sampler_cfg['train_for_val'] = pop_without_del(sampler_cfg, 'train_for_val', None)
        cut_n_images = sampler_cfg['train_for_val'].n_images or len(train_dataset)
        train_for_val_sampler = train_sampler.copy(sequential_override=True,
                                                   cut_n_images=None if cut_n_images is None
                                                   else min(cut_n_images, len(train_sampler)))

    return train_sampler, val_sampler, train_for_val_sampler

def sampler_generator_helper(dataset_type, dataset, default_dataset, sampler_cfg, sampler_type, transformer_tag):
    instance_count_file = os.path.join(dataset_registry.REGISTRY[dataset_type].dataset_path,
                                       '{}_instance_counts_{}.npy'.format(sampler_type, transformer_tag))
    sem_cls_filter = sampler_cfg[sampler_type].sem_cls_filter
    sem_cls_filter_values = convert_sem_cls_filter_from_names_to_values(
        sem_cls_filter, default_dataset.semantic_class_names) \
        if sem_cls_filter is not None and isinstance(sem_cls_filter[0], str) else sem_cls_filter
    sampler = get_configured_sampler(
        dataset, default_dataset, sequential=True,
        n_instances_range=sampler_cfg[sampler_type].n_instances_range,
        n_images=sampler_cfg[sampler_type].n_images,
        sem_cls_filter=sem_cls_filter_values, instance_count_file=instance_count_file)
    return sampler

def convert_sem_cls_filter_from_names_to_values(sem_cls_filter, semantic_class_names):
    try:
        sem_cls_filter_values = [semantic_class_names.index(class_name)
                                 for class_name in sem_cls_filter]
    except:
        sem_cls_filter_values = [int(np.where(semantic_class_names == class_name)[0][0])
                                 for class_name in sem_cls_filter]
    return sem_cls_filter_values


import pygame
from pygame.sprite import Sprite

class Power(Sprite):
    def __init__(self, center, screen):
        super(Power, self).__init__()
        self.screen = screen
        self.image = pygame.image.load('images/heart.png')
        self.rect = self.image.get_rect()
        #self.rect = pygame.Rect(0, 0, 5, 30)
        self.rect.center = center

        self.speed_y = 2

    def blitme(self):
        #pygame.draw.rect(self.screen, self.color, self.rect)
        self.screen.blit(self.image, self.rect)

    def update(self):
        self.rect.y += self.speed_y
        if self.rect.top > 600:
            self.kill()


import re
import six
from collections import OrderedDict

from src.db_parser.tokens import TokenTypes
from src.db_parser.common import DbSyntaxError

class Token(object):
    """
    Class representing a lexer token. Tokens are considered equal if their types are equal.
    Args:
        type: the type of this token. Should be one of TokenTypes
        linenum: the line number this token was found on
        colnum: the column number this token was found on
        contents: the original text that this token was parsed from
    """

    def __init__(self, type, linenum, colnum, contents=None):
        self.type = type
        self.contents = contents

        self.line = linenum
        self.col = colnum

    def __str__(self):
        return "{} (contents={})".format(self.type, self.contents)

    __repr__ = __str__

    def __eq__(self, other):
        try:
            return self.type == other.type
        except AttributeError:
            return False

def _escape(var):
    """
    Returns the input variable escaped and wrapped in a regex capture group.
    """
    return f"({re.escape(var)})"

class Lexer(six.Iterator):
    """
    Lexer, tokenises the database file into
    """

    """
    Tokens to ignore. These are never returned from __next__
    """
    IGNORED_TOKENS = [TokenTypes.WHITESPACE]

    """
    This provides a mapping between regexes and lexer tokens.

    The regexes are tried in order, the first one that matches get's it's token produced.
    """
    TOKEN_MAPPING = OrderedDict([
        (_escape("$("),  # Macro start with brackets $(MACRO=VALUE)
         TokenTypes.BRACKET_MACRO_START),
        (_escape("${"),  # Macro start with curly braces ${MACRO=VALUE}
         TokenTypes.BRACE_MACRO_START),
        (_escape("record"),
         TokenTypes.RECORD),
        (_escape("grecord"),
         TokenTypes.RECORD),
        (_escape("field"),
         TokenTypes.FIELD),
        (_escape("info"),
         TokenTypes.INFO),
        (_escape("alias"),
         TokenTypes.ALIAS),
        (_escape("("),
         TokenTypes.L_BRACKET),
        (_escape(")"),
         TokenTypes.R_BRACKET),
        (_escape("{"),
         TokenTypes.L_BRACE),
        (_escape("}"),
         TokenTypes.R_BRACE),
        (_escape(","),
         TokenTypes.COMMA),
        (_escape("="),
         TokenTypes.EQUALS),
        (r"(\".*?[^\\]\"|\"\")",  # Ignore escaped quotes within a string. Add special case for empty string
         TokenTypes.QUOTED_STRING),
        (r"(\s+)",
         TokenTypes.WHITESPACE),
        (r"([a-zA-Z0-9\-\_\.\:]+)",  # Alphanumeric, -, _, ., :
         TokenTypes.LITERAL),
        (_escape("#"),  # Mostly comments, but also things like $(MACRO=#) field(INPA, "blah") are valid...
         TokenTypes.HASH),
        (r"(.+)",  # Matches absolutely anything (as a last resort, for comment contents for example).
         TokenTypes.UNKNOWN),
    ])

    # Change to keep track of macro
    def __init__(self, file_contents):
        self.file_contents = file_contents
        self.gen = None

    def token_generator(self):
        """
        Token generator function.
        yields:
            Tokens corresponding to the lexed input.
        """
        lines = self.file_contents.split("\n")
        for linenum, line in enumerate(lines, 1):
            column = 0
            while column < len(line):
                for regexp in Lexer.TOKEN_MAPPING:
                    match_text = self._text_matches_regex(line[column:], regexp)
                    if match_text is not None:
                        yield Token(Lexer.TOKEN_MAPPING[regexp], linenum, column, match_text)
                        column += len(match_text)
                        break
                else:
                    raise DbSyntaxError("No matching rules found at {}:{}. Line contents: '{}'"
                                        .format(linenum, column, line))

        yield Token(TokenTypes.EOF, len(lines), len(lines[len(lines) - 1]))

    def __next__(self):
        """
        Generator of this Lexer's tokens
        """
        if self.gen is None:
            self.gen = self.token_generator()

        tok = next(self.gen)
        while tok.type in self.IGNORED_TOKENS:
            tok = next(self.gen)
        return tok

    @staticmethod
    def _text_matches_regex(text, rx):
        """
        Attempts to match the given text with a given regex.

        Args:
            text:
            rx: the regex to apply
        Returns:
             The text of the match if the match succeeded, None otherwise
        """
        match = re.match(rx, text)
        if match:
            return match.group(1)
        return None


from sklearn.cluster import KMeans
from skimage.io import imread
from skimage import img_as_float as iaf
from skimage.io import imshow
import numpy as np
from math import log10

n_colors = 20  #clusters

image   = iaf(imread('parrots.jpg'))

'''features matrix'''
r   = np.array([image[:,:,0].ravel()]).T #T - transpose
g   = np.array([image[:,:,1].ravel()]).T
b   = np.array([image[:,:,2].ravel()]).T
X   = np.hstack((r,g,b))

w, h, d = image.shape                    #matrix sizes

'''clusterization'''
clr = KMeans(n_clusters = 8, init = 'k-means++', random_state=241)
clr.fit(X)

def recreate_image(codebook, labels, w, h, d):
    """Recreate the (compressed) image from the code book & labels"""
    image = np.zeros((w, h, d))
    label_idx = 0
    for i in range(w):
        for j in range(h):
            image[i][j] = codebook[labels[label_idx]]
            label_idx += 1
    return image

new_image = recreate_image(clr.cluster_centers_, clr.labels_, w, h, d)

imshow(new_image)

mse = np.mean((image - new_image) ** 2)
psnr = 10 * log10(float(1) / mse)

print(psnr) #psnr > 20 if n_clusters = 11


##############################################################################
#   Copyright (c) 2018 Eynes/E-MIPS (www.eynes.com.ar)
#   License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
##############################################################################

from lxml import etree
from datetime import datetime
import time
import email
import urllib.error as Urllib
from xml.sax import SAXParseException
import logging
import pytz

# Configuracion del logger
logger = logging.getLogger('afipws')
logger.setLevel(logging.DEBUG)
try:
    from suds.client import Client
except (ImportError, IOError) as e:
    logger.debug("Cannot import Client from Suds-py3:\n%s" % e)
try:
    from M2Crypto import BIO, SMIME
except (ImportError, IOError) as e:
    logger.debug("Cannot import BIO & SMIME from M2Crypto: \n%s" % e)

class WSAA:
    def __init__(self, cert, private_key, wsaaurl, service, tz):
        """Init."""
        self.token = None
        self.sign = None
        self.wsaaurl = wsaaurl
        self.service = service
        self.expiration_time = None
        self.ta = None
        self.connected = True
        self.timezone = tz

        try:
            self._create_client()
        except Urllib.URLError as e:
            logger.error("No hay conexion disponible")
            self.connected = False
            raise Exception('No se puede conectar con AFIP: %s' % str(e))
        except SAXParseException as e:
            raise Exception('WSAA URL malformada: %s' % e.getMessage())
        except Exception as e:
            raise Exception('Unknown Error: %s' % e)

    def _create_client(self):
        # Creamos el cliente contra la URL
        self.client = Client(self.wsaaurl)

    def _create_tra(self):
        """Funcion para crear el TRA."""
        # Creamos el root
        logger.debug("Creando TRA")
        root = etree.Element('loginTicketRequest')
        doc = etree.ElementTree(root)
        root.set('version', '1.0')

        # Creamos los childs de loginTicketRequest
        header = etree.SubElement(root, 'header')

        # UniqueId
        uniqueid = etree.SubElement(header, 'uniqueId')
        timestamp = int(time.mktime(datetime.now().timetuple()))
        uniqueid.text = str(timestamp)

        # generationTime
        tsgen = datetime.fromtimestamp(timestamp-2400)
        tsgen = pytz.utc.localize(tsgen).astimezone(self.timezone)
        gentime = etree.SubElement(header, 'generationTime')
        gentime.text = tsgen.isoformat()

        # expirationTime
        tsexp = datetime.fromtimestamp(timestamp+14400)
        tsexp = pytz.utc.localize(tsexp).astimezone(self.timezone)
        exptime = etree.SubElement(header, 'expirationTime')
        exptime.text = tsexp.isoformat()

        # service
        serv = etree.SubElement(root, 'service')
        serv.text = self.service

        # try:
        #     f = open('tra.xml', 'w')
        #     doc.write(f, xml_declaration=True, encoding='UTF-8',
        #               pretty_print=True)
        # except Exception, e:
        #     print 'No se puede grabar el archivo: %s' % e
        #     return None

        logger.debug("TRA Creado")
        return etree.tostring(doc)

    def _sign_tra(self, tra, cert, key):
        """Funcion que firma el tra."""

        # Creamos un buffer a partir del TRA
        buf = BIO.MemoryBuffer(tra)
        key_bio = BIO.MemoryBuffer(key.encode('ascii'))
        cert_bio = BIO.MemoryBuffer(cert.encode('ascii'))

        # Firmamos el TRA
        s = SMIME.SMIME()
        s.load_key_bio(key_bio, cert_bio)
        p7 = s.sign(buf, 0)
        out = BIO.MemoryBuffer()
        s.write(out, p7)

        # Extraemos la parte que nos interesa
        msg = email.message_from_string(out.read().decode('utf-8'))
        for part in msg.walk():
            filename = part.get_filename()
            if filename == "smime.p7m":
                logger.debug("TRA Firmado")
                return part.get_payload(decode=False)

    def _call_wsaa(self, cms):
        """ Funcion para llamar al WSAA y loguearse """

        # Si no esta conectado, probamos conectar nuevamente
        if not self.connected:
            try:
                self._create_client()
            except Urllib.URLError as e:
                logger.warning("No hay conexion disponible")
                self.connected = False
                raise Exception('No hay conexion: %s' % e)

        # Llamamos a loginCms y obtenemos el resultado
        logger.debug("Llamando a loginCms:\n%s", cms)
        try:
            result = self.client.service.loginCms(cms)
        except Exception as e:
            logger.exception("Excepcion al llamar a loginCms")
            raise Exception('Exception al autenticar: %s' % str(e))

        self.ta = result
        return result

    def parse_ta(self, ta=None):
        if not ta:
            ta = self.ta
            if not ta:
                return

        # Quitamos la declaracion de XML
        tas = ta.split('\n')
        if tas[0].find("?xml"):
            ta = '\n'.join(tas[1:])

        # Parseamos el resultado
        root = etree.XML(ta)

        # Buscamos el expirationTime
        header = root.find('header')
        exptime = header.find('expirationTime')
        exptime_iso = exptime.text[0:exptime.text.rfind('-')]
        self.expiration_time = datetime.strptime(exptime_iso,
                                                 '%Y-%m-%dT%H:%M:%S.%f')

        # Si no expiro, obtenemos el token y el sign
        cred = root.find('credentials')
        self.token = cred.find('token').text
        self.sign = cred.find('sign').text

        return True

    # TODO: Agregar una flag de force para tomarlo igual
    # TODO: Hacer chequeo de errores
    def get_token_and_sign(self, cert, key, force=True):
        # Primero chequeamos si ya tenemos un token
        if not force:
            if self.ta and self.expiration_time and self.token and self.sign:
                # Si todavia no expiro el que tenemos, lo retornamos
                if datetime.now() < self.expiration_time:
                    return self.token, self.sign

        tra = self._create_tra()
        cms = self._sign_tra(tra, cert, key)
        try:
            self._call_wsaa(cms)
        except Exception as e:
            raise e

        self.parse_ta(self.ta)
        return True

# El certificado X.509 obtenido de Seg. Inf.
# CERT = "/home/skennedy/proyectos/afipws/certs2012/eynes/cert_eynes.crt"
CERT = ""
# La clave privada del certificado CERT
# PRIVATEKEY = "/home/skennedy/proyectos/afipws/certs2012/eynes/ \
#     privada_eynes.key"
PRIVATEKEY = ""
# homologacion (pruebas)
# WSAAURL = "https://wsaahomo.afip.gov.ar/ws/services/LoginCms?wsdl"
WSAAURL = ""
# WSAAURL="https://wsaa.afip.gov.ar/ws/services/LoginCms?wsdl"

if __name__ == '__main__':
    wsaa = WSAA(CERT, PRIVATEKEY, WSAAURL, "wsfe")
    token = None
    sign = None
    expiration_time = None
    try:
        # Se puede llamar a get_token_and_sign o pasarle un TA
        # anterior a wsaa.parse_ta
        wsaa.get_token_and_sign()
        token = wsaa.token
        expiration_time = wsaa.expiration_time
        sign = wsaa.sign
    except Exception as e:
        print(e)

    # Vemos si ya expiro
    if datetime.now() > expiration_time:
        print("Certificado expirado")

    print('Token: ', token)
    print('Sign: ', sign)
    print('Expiration Time: ', expiration_time.strftime("%d/%m/%Y %H:%M:%S"))


# coding=utf-8
# Q18:
# A hapax legomenon (often abbreviated to hapax) is a word which occurs only once in either the written record
# of a language, the works of an author, or in a single text. Define a function that given the file name
# of a text will return all its hapaxes. Make sure your program ignores capitalization.

import operator
from collections import Counter

def hapax(filename):
    f = open(filename)
    try:
        counter = Counter(f.read().split())
    except (Exception,) as e:
        print(e)
    else:
        for k, v in sorted(counter.iteritems(), key=operator.itemgetter(0)):
            if v == 1 and len(k) > 3:
                print(k)

    finally:
        f.close()

hapax(__file__)


import sys
import argparse
import glob
from yolo import YOLO, detect_video
from PIL import Image

count = 0

def person_exist_in_obj(obj):
    i = 0
    while True:
        try:
            if obj[i]["class"] == "person":
                return True
            i += 1
        except:
            break
    return False

def crop_person(image, obj):
    global count
    i = 0
    while True:
        try:
            if obj[i]["class"] == "person":
                fname = './crop_person_images/' + str(count) + '.jpg'
                img_crop = image.crop(obj[i]["bbox"])
                img_crop.save(fname)
                count += 1
                # img_crop.show()
            i += 1
        except:
            break
    return False

def detect_img(yolo):
    image_path = glob.glob('./mscoco2017/train2017/*.jpg')
    i = 0
    while True:
        try:
            img = image_path[i]
            print(img)
            i += 1
        except:
            break
        try:
            image = Image.open(img)
            original_image = image.copy()
        except:
            print('Open Error! Try again!')
            continue
        else:
            try:
                r_image, obj = yolo.detect_image(image)
            except:
                continue
            if person_exist_in_obj(obj):
                crop_person(original_image, obj)

            # r_image.show()
    yolo.close_session()

FLAGS = None

if __name__ == '__main__':
    # class YOLO defines the default value, so suppress any default here
    parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
    '''
    Command line options
    '''

    parser.add_argument(
        '-m', '--model', type=str,
        help='path to model weight file'
    )

    parser.add_argument(
        '-a', '--anchors', type=str,
        help='path to anchor definitions'
    )

    parser.add_argument(
        "-c", '--classes', type=str,
        help='path to class definitions'
    )

    parser.add_argument(
        "-g", '--gpu_num', type=int, default=1,
        help='Number of GPU to use'
    )

    parser.add_argument(
        '--image', default=False, action="store_true",
        help='Image detection mode, will ignore all positional arguments'
    )
    '''
    Command line positional arguments -- for video detection mode
    '''
    parser.add_argument(
        "--input", nargs='?', type=str,required=False,default=0,
        help = "Video input path"
    )

    parser.add_argument(
        "--output", nargs='?', type=str, default="",
        help = "[Optional] Video output path"
    )

    FLAGS = parser.parse_args()

    if FLAGS.image:
        """
        Image detection mode, disregard any remaining command line arguments
        """
        print("Image detection mode")
        detect_img(YOLO(**vars(FLAGS)))
    elif "input" in FLAGS:
        detect_video(YOLO(**vars(FLAGS)), FLAGS.input, FLAGS.output)
    else:
        print("Must specify at least video_input_path.  See usage with --help.")


from django.http import (HttpResponse, HttpResponseRedirect)
from django.shortcuts import render
from django.core.urlresolvers import reverse
from django.utils.datastructures import MultiValueDictKeyError
from django.shortcuts import get_object_or_404
from django.contrib.auth.decorators import login_required
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
from django.conf import settings
from django.forms.models import save_instance

from whats_fresh.whats_fresh_api.models import (Product, Preparation,
                                                ProductPreparation)
from whats_fresh.whats_fresh_api.forms import ProductForm
from whats_fresh.whats_fresh_api.functions import group_required
from whats_fresh.whats_fresh_api.views.serializer import FreshSerializer
from whats_fresh.whats_fresh_api.templatetags import get_fieldname

from haystack.query import SearchQuerySet
from collections import OrderedDict
import json

@login_required
@group_required('Administration Users', 'Data Entry Users')
def product(request, id=None):
    """
    */entry/products/<id>*, */entry/products/new*

    The entry interface's edit/add/delete product view. This view creates the
    edit page for a given product, or the "new product" page if it is not
    passed an ID. It also accepts POST requests to create or edit products, and
    DELETE requests to delete them.

    If called with DELETE, it will return a 200 upon success or a 404 upon
    failure. This is to be used as part of an AJAX call, or some other API
    call.
    """
    if request.method == 'DELETE':
        product = get_object_or_404(Product, pk=id)
        product.delete()
        return HttpResponse()

    if request.method == 'POST':
        message = ''
        post_data = request.POST.copy()
        errors = []

        try:
            if len(post_data['preparation_ids']) == 0:
                errors.append(
                    "You must choose at least one entry from "
                    + get_fieldname.get_fieldname('preparations'))
                preparations = []
            else:
                preparations = [int(p) for p in set(
                    post_data['preparation_ids'].split(','))]
        except MultiValueDictKeyError:
            errors.append(
                "You must choose at least one entry from "
                + get_fieldname.get_fieldname('preparations'))
            preparations = []

        if id:
            product = Product.objects.get(id=id)
        else:
            product = None

        product_form = ProductForm(post_data, product)
        if product_form.is_valid() and not errors:
            if id:
                for product_preparation in (
                        ProductPreparation.objects.filter(product=product)
                        .exclude(preparation__id__in=preparations)):
                    product_preparation.delete()

                existing_preparations = [
                    pp.preparation.id for pp in
                    ProductPreparation.objects.filter(
                        product=product,
                        preparation__id__in=preparations)
                ]

                for preparation in preparations:
                    if preparation not in existing_preparations:
                        ProductPreparation.objects.create(
                            product=product,
                            preparation=Preparation.objects.get(
                                id=preparation))

                save_instance(product_form, product)
            else:
                product = Product.objects.create(**product_form.cleaned_data)
                for preparation in preparations:
                    product_preparation = ProductPreparation.objects.create(
                        product=product,
                        preparation=Preparation.objects.get(
                            id=preparation))
                product.save()
            return HttpResponseRedirect(
                "%s?saved=true" % reverse('entry-list-products'))
    else:
        errors = []
        message = ''

    if id:
        product = Product.objects.get(id=id)
        title = "Edit {0}".format(product.name)
        post_url = reverse('edit-product', kwargs={'id': id})
        product_form = ProductForm(instance=product)

        existing_preparations = product.preparations.all()

        if request.GET.get('success') == 'true':
            message = "Product saved successfully!"

    elif request.method != 'POST':
        product_form = ProductForm()
        post_url = reverse('new-product')
        title = "New " + get_fieldname.get_fieldname('products')
        existing_preparations = []

    else:
        post_url = reverse('new-product')
        title = "New " + get_fieldname.get_fieldname('products')
        existing_preparations = []

    data = {'preparations': []}

    for preparation in Preparation.objects.all():
        data['preparations'].append({
            'id': preparation.id,
            'name': preparation.name
        })

    json_preparations = json.dumps(data)

    return render(request, 'product.html', {
        'parent_url': [
            {'url': reverse('home'), 'name': 'Home'},
            {'url': reverse('entry-list-products'),
             'name': get_fieldname.get_fieldname('products')}],
        'json_preparations': json_preparations,
        'preparation_dict': data,
        'existing_preparations': existing_preparations,
        'parent_text': 'Product List',
        'message': message,
        'title': title,
        'post_url': post_url,
        'errors': errors,
        'product_form': product_form,
    })

@login_required
@group_required('Administration Users', 'Data Entry Users')
def product_list(request):
    """
    */entry/products*

    The entry interface's products list. This view lists all products,
    their description, and allows you to click on them to view/edit the
    product.
    """

    message = ""
    if request.GET.get('success') == 'true':
        message = "Entry deleted successfully!"
    elif request.GET.get('saved') == 'true':
        message = "Entry saved successfully!"

    search = request.GET.get('search')

    if search is None or search.strip() == "":
        products = Product.objects.order_by('name')
    else:
        if request.GET.get('search') != "":
            products = list(
                OrderedDict.fromkeys(
                    item.object for item in
                    SearchQuerySet().models(Product).autocomplete(
                        content_auto=search)))
            if not products:
                message = "No results"

    paginator = Paginator(products, settings.PAGE_LENGTH)

    page = request.GET.get('page')

    try:
        products = paginator.page(page)
    except PageNotAnInteger:
        # If page is not an integer, deliver first page.
        products = paginator.page(1)
    except EmptyPage:
        # If page is out of range (e.g. 9999), deliver last page of results.
        products = paginator.page(paginator.num_pages)

    return render(request, 'list.html', {
        'message': message,
        'parent_url': reverse('home'),
        'parent_text': 'Home',
        'new_url': reverse('new-product'),
        'title': get_fieldname.get_fieldname('products'),
        'item_list': products,
        'description_field': {'title': 'Variety', 'attribute': 'variety'},
        'edit_url': 'edit-product',
        'search_text': request.GET.get('search'),
        'list_url': get_fieldname.get_fieldname('products_slug')

    })

@login_required
@group_required('Administration Users', 'Data Entry Users')
def product_ajax(request, id=None):
    if request.method == 'GET':
        product_form = ProductForm()

        data = {'preparations': []}

        for preparation in Preparation.objects.all():
            data['preparations'].append({
                'id': preparation.id,
                'name': preparation.name
            })

        json_preparations = json.dumps(data)

        return render(request, 'product_ajax.html',
                      {'product_form': product_form,
                       'json_preparations': json_preparations,
                       'preparation_dict': data})

    elif request.method == 'POST':
        message = ''
        post_data = request.POST.copy()
        errors = []

        try:
            if len(post_data['prep_ids']) == 0:
                errors.append(
                    "You must choose at least one entry from "
                    + get_fieldname.get_fieldname('preparations'))
                preparations = []
            else:
                preparations = [int(p) for p in set(
                    post_data['prep_ids'].split(','))]
        except MultiValueDictKeyError:
            errors.append("You must choose at least one entry from " +
                          get_fieldname.get_fieldname('preparations'))
            preparations = []

        product = None

        product_form = ProductForm(post_data, product)
        popup_prep = []
        if product_form.is_valid() and not errors:
            product = Product.objects.create(**product_form.cleaned_data)
            for preparation in preparations:
                inline_prep = ProductPreparation.objects.create(
                    product=product,
                    preparation=Preparation.objects.get(
                        id=preparation))
                popup_prep.append({
                    'id': inline_prep.id,
                    'name': inline_prep.preparation.name
                })
                product.save()
            serializer = FreshSerializer()
            data = json.loads(serializer.serialize(product))
            data['preparations'] = popup_prep
            return HttpResponse(json.dumps(data),
                                content_type="application/json")

        data = {'preparations': []}

        for preparation in Preparation.objects.all():
            data['preparations'].append({
                'id': preparation.id,
                'name': preparation.name})

        json_preparations = json.dumps(data)

        return render(request, 'product_ajax.html', {
            'json_preparations': json_preparations,
            'preparation_dict': data,
            'parent_text': 'Product List',
            'message': message,
            'errors': errors,
            'product_form': product_form})


# vi: ts=4 expandtab
#
#    Author: Edouard Bonlieu <ebonlieu@ocs.online.net>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License version 3, as
#    published by the Free Software Foundation.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

import functools
import json

from requests.packages.urllib3.poolmanager import PoolManager
import requests

from cloudinit import log as logging
from cloudinit import sources
from cloudinit import url_helper
from cloudinit import util

LOG = logging.getLogger(__name__)

BUILTIN_DS_CONFIG = {
    'metadata_url': 'http://169.254.42.42/conf?format=json',
    'userdata_url': 'http://169.254.42.42/user_data/cloud-init'
}

DEF_MD_RETRIES = 5
DEF_MD_TIMEOUT = 10

class SourceAddressAdapter(requests.adapters.HTTPAdapter):
    """ Adapter for requests to choose the local address to bind to.
    """

    def __init__(self, source_address, **kwargs):
        self.source_address = source_address
        super(SourceAddressAdapter, self).__init__(**kwargs)

    def init_poolmanager(self, connections, maxsize, block=False):
        self.poolmanager = PoolManager(num_pools=connections,
                                       maxsize=maxsize,
                                       block=block,
                                       source_address=self.source_address)

class DataSourceScaleway(sources.DataSource):

    def __init__(self, sys_cfg, distro, paths):
        LOG.debug('Init scw')
        sources.DataSource.__init__(self, sys_cfg, distro, paths)

        self.metadata = {}
        self.ds_cfg = util.mergemanydict([
            util.get_cfg_by_path(sys_cfg, ["datasource", "Scaleway"], {}),
            BUILTIN_DS_CONFIG
        ])

        self.metadata_address = self.ds_cfg['metadata_url']
        self.userdata_address = self.ds_cfg['userdata_url']

        self.retries = self.ds_cfg.get('retries', DEF_MD_RETRIES)
        self.timeout = self.ds_cfg.get('timeout', DEF_MD_TIMEOUT)

    def _get_metadata(self):
        resp = url_helper.readurl(
            self.metadata_address,
            data=None,
            timeout=self.timeout,
            retries=self.retries
        )
        metadata = json.loads(util.decode_binary(resp.contents))
        LOG.debug('metadata downloaded')

        # try to make a request on the first privileged port available
        for port in range(1, 1024):
            try:
                session = requests.Session()
                session.mount(
                    'http://',
                    SourceAddressAdapter(source_address=('0.0.0.0', port))
                )
                resp = url_helper.readurl(
                    self.userdata_address,
                    data=None,
                    timeout=self.timeout,
                    retries=self.retries,
                    session=session
                )
                user_data = util.decode_binary(resp.contents)
                LOG.debug('user-data downloaded')
                return metadata, user_data

            except url_helper.UrlError:  # try next port
                pass

    def get_data(self):
        metadata, metadata['user-data'] = self._get_metadata()
        self.metadata = {
            'id': metadata['id'],
            'hostname': metadata['hostname'],
            'user-data': metadata['user-data'],
            'ssh_public_keys': [
                key['key'] for key in metadata['ssh_public_keys']
            ]
        }
        return True

    @property
    def launch_index(self):
        return None

    def get_instance_id(self):
        return self.metadata['id']

    def get_public_ssh_keys(self):
        return self.metadata['ssh_public_keys']

    def get_hostname(self, fqdn=False, resolve_ip=False):
        return self.metadata['hostname']

    def get_userdata_raw(self):
        return self.metadata['user-data']

    @property
    def availability_zone(self):
        return None

    @property
    def region(self):
        return None

datasources = [
    (DataSourceScaleway, (sources.DEP_FILESYSTEM, sources.DEP_NETWORK)),
]

def get_datasource_list(depends):
    return sources.list_from_depends(depends, datasources)


"""Utility Widgets"""
from itertools import groupby

from django.forms.models import ModelChoiceField, ModelChoiceIterator

class GroupedModelChoiceIterator(ModelChoiceIterator):
    """ModelChoiceField which groups objects by their verbose_name"""

    def __iter__(self):
        if self.field.empty_label is not None:
            yield ("", self.field.empty_label)
        queryset = self.queryset
        # Can't use iterator() when queryset uses prefetch_related()
        if not queryset._prefetch_related_lookups:
            queryset = queryset.iterator()
        # We can't use DB-level sorting as we sort by subclass
        queryset = sorted(queryset, key=lambda x: x._meta.verbose_name)
        for group, objs in groupby(queryset, key=lambda x: x._meta.verbose_name):
            yield (group, [self.choice(obj) for obj in objs])

class GroupedModelChoiceField(ModelChoiceField):
    """ModelChoiceField which groups objects by their verbose_name"""

    iterator = GroupedModelChoiceIterator


import numpy as np
import torch
import torch.nn as nn

from rotation_utils import transform3d
from utils import transform_voxel_to_match_image

class Generator(nn.Module):

    def __init__(self, args, c_dim=3, gf_dim=64, voxel_size=128):

        super().__init__()
        self.image_size = args.image_size
        self.batch_size = args.batch_size

        self.gf_dim = gf_dim
        self.voxel_size = voxel_size
        self.c_dim = c_dim  # out_channels

        self.in_channels = self.gf_dim * 8  # 512

        self.z_mapping_mlp = ZMappingMLP()
        self.adain = AdaIn()

        # in_channels => const_tensor.shape[1]
        self.deconv3d_1 = nn.ConvTranspose3d(in_channels=self.in_channels, out_channels=self.gf_dim * 2, kernel_size=2, stride=2, padding=0)
        self.deconv3d_2 = nn.ConvTranspose3d(in_channels=self.gf_dim * 2, out_channels=self.gf_dim * 1, kernel_size=2, stride=2, padding=0)

        # convolution 2d
        # for projection
        self.deconv2d_1 = nn.ConvTranspose2d(in_channels=self.gf_dim * self.voxel_size, out_channels=self.gf_dim * 16, kernel_size=1, stride=1, padding=0)

        self.deconv2d_2 = nn.ConvTranspose2d(in_channels=self.gf_dim * 16, out_channels=self.gf_dim * 4, kernel_size=3, stride=1, padding=1)
        self.deconv2d_3 = nn.ConvTranspose2d(in_channels=self.gf_dim * 4, out_channels=self.gf_dim, kernel_size=3, stride=1, padding=1)

        self.deconv2d_4 = nn.ConvTranspose2d(in_channels=self.gf_dim, out_channels=self.c_dim, kernel_size=3, stride=1, padding=1)

    def forward(self, x, z, view_in):
        """
        input => x (batch_size, channel(512), height(4), width(4), depth(4)), z => (cont_dim vector)
        conv3dblock * 2 => conv3d, adain(args: MLP(z)), leakyrelu

        3d transformer(args: camera pose)

        conv3d block * 2

        projection unit

        conv2d block => conv2d, adain(args: MLP(z)), leakyrelu

        """
        # Input block
        s0, b0 = self.z_mapping_mlp(z, self.in_channels)  # s => scale, b => bias, adain = scale * ((x - mean) / std) + bias
        h0 = self.adain(x, s0, b0)
        h0 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(h0)

        # print('first hidden shape:', h0.shape)  # (batch_size, 512, 4, 4, 4)

        # Block1
        h1 = self.deconv3d_1(h0)
        s1, b1 = self.z_mapping_mlp(z, self.gf_dim * 2)
        h1 = self.adain(h1, s1, b1)
        h1 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(h1)  # (batch_size, 128, 8, 8, 8)

        # Block2
        h2 = self.deconv3d_2(h1)
        s2, b2 = self.z_mapping_mlp(z, self.gf_dim * 1)
        h2 = self.adain(h2, s2, b2)
        h2 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(h2)  # (batch_size, 64, 16, 16, 16)

        # 3d transformer(args view_in(camera pose))
        # view_in => 6 dimension tensor
        h2_rotated = transform3d(voxel_array=h2, view_params=view_in)  # output => voxel format(b, 1, 128, 128, 128)

        # projection unit(with depth dimension collapsion)
        h2_rotated = transform_voxel_to_match_image(h2_rotated)

        batch_size, channels, depth, height, width = h2_rotated.shape

        # collapsing depth dimension  (batch_size, channels, depth, height, width)
        h2_2d = h2_rotated.reshape(batch_size, channels * depth, height, width)  # (batch_size, channels * depth, height, width)

        # 1 * 1 convolution(occulusion learning for collapsing dimension)
        h3 = self.deconv2d_1(h2_2d)
        h3 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(h3)

        # 2d convolution block
        h4 = self.deconv2d_2(h3)
        s4, b4 = self.z_mapping_mlp(z, self.gf_dim * 4)  # ToDo: same function is ok?
        h4 = self.adain(h4, s4, b4)
        h4 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(h4)

        h5 = self.deconv2d_3(h4)
        s5, b5 = self.z_mapping_mlp(z, self.gf_dim)
        h5 = self.adain(h5, s5, b5)
        h5 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(h5)
        h6 = self.deconv2d_4(h5)
        output = nn.Tanh()(h6)

        return output

class Discriminator(nn.Module):

    def __init__(self, args, in_channels=3, df_dim=64):

        super().__init__()

        self.df_dim = df_dim

        # input size => (batch_size, 3, 128, 128)
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=self.df_dim, kernel_size=3, stride=2, padding=1)  # (batch_size, 64, 64, 64)
        self.conv2 = nn.Conv2d(in_channels=self.df_dim, out_channels=self.df_dim * 2, kernel_size=3, stride=2, padding=1)  # (batch_size, 128, 32, 32)
        self.instance_norm2 = nn.InstanceNorm2d(num_features=self.df_dim * 2)
        self.conv3 = nn.Conv2d(in_channels=self.df_dim * 2, out_channels=self.df_dim * 4, kernel_size=3, stride=2, padding=1)  # (batch_size, 256, 16, 16)
        self.instance_norm3 = nn.InstanceNorm2d(num_features=self.df_dim * 4)
        self.conv4 = nn.Conv2d(in_channels=self.df_dim * 4, out_channels=self.df_dim * 8, kernel_size=3, stride=2, padding=1)  # (batch_size, 512, 8, 8)
        self.instance_norm4 = nn.InstanceNorm2d(num_features=self.df_dim * 8)

        # additional layer
        flatten_size = int(self.df_dim * 8 * np.floor(args.image_size // 2 ** 4) ** 2)
        self.linear1 = nn.Linear(in_features=flatten_size, out_features=1)

        self.linear2 = nn.Linear(in_features=flatten_size, out_features=128)
        self.linear3 = nn.Linear(in_features=128, out_features=args.z_dim)

    def forward(self, x):
        """

        h0 => leakyrelu(conv(x))
        h1 => leakyrelu(instance_norm(conv_specnorm(h0)))
        h2 => leakyrelu(instance_norm(conv_specnorm(h1)))
        h3 => leakyrelu(instance_norm(conv_specnorm(h2)))
        h3 => flatten()
        h4 => linear(h3)

        # additional layer
        encoder => leakyrelu(linear(flatten(h3), 128))
        output => linear(encoder, cont_dim)

        """

        # recognition network for latent variables has an additional layer
        # linear layer
        # output => cont_dim vector
        batch_size = x.shape[0]
        h0 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(self.conv1(x))
        h1 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(self.instance_norm2(self.conv2(h0)))
        h2 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(self.instance_norm3(self.conv3(h1)))
        h3 = nn.LeakyReLU(negative_slope=0.2, inplace=True)(self.instance_norm4(self.conv4(h2)))
        flatten_h3 = h3.view(batch_size, -1)

        # additional layer
        h4 = self.linear1(flatten_h3)

        encoder = nn.LeakyReLU(negative_slope=0.2, inplace=True)(self.linear2(flatten_h3))
        latent_output = nn.Tanh()(self.linear3(encoder))

        return h4, nn.Sigmoid()(h4), latent_output, (nn.Sigmoid()(h1), nn.Sigmoid()(h2), nn.Sigmoid()(h3), nn.Sigmoid()(h4))

class ZMappingMLP(nn.Module):

    def __init__(self):

        super().__init__()

    def forward(self, x, out_channels):
        # scale = sigma(y), bias = mu(y)
        # z = (batch_size, z_dim)
        x = nn.Linear(in_features=x.shape[-1], out_features=out_channels * 2)(x)
        x = nn.LeakyReLU(negative_slope=0.2, inplace=True)(x)
        scale = x[:, :out_channels]
        bias = x[:, out_channels:]

        return scale, bias

class AdaIn(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x, s, b):

        size = x.size()
        batch_size, in_channels = s.shape

        if len(x.shape) == 5:
            mean = torch.mean(x.view(batch_size, in_channels, -1), 2).view(batch_size, in_channels, 1, 1, 1)
            std = torch.std(x.view(batch_size, in_channels, -1), 2).view(batch_size, in_channels, 1, 1, 1)
            s = s.view(batch_size, in_channels, 1, 1, 1).expand(size)
            b = b.view(batch_size, in_channels, 1, 1, 1).expand(size)
        else:
            mean = torch.mean(x.view(batch_size, in_channels, -1), 2).view(batch_size, in_channels, 1, 1)
            std = torch.std(x.view(batch_size, in_channels, -1), 2).view(batch_size, in_channels, 1, 1)
            s = s.view(batch_size, in_channels, 1, 1).expand(size)
            b = b.view(batch_size, in_channels, 1, 1).expand(size)

        # match dimension (batch_size, channles, depth, height, width)
        mean = mean.expand(size)
        std = std.expand(size)

        # print('scale shape:', s.shape, 'bias shape:', b.shape, 'mean:', mean.shape, 'std:', std.shape)

        output = s * (x - mean) / std + b
        return output


import matplotlib as mpl
mpl.use('agg')
from lauren import *
from radial_data import *
import matplotlib.pyplot as plt

#fn="/Users/laurennc/data/RyanSims/r0058_l10/redshift0058"
fn="/hpc/astrostats/astro/users/lnc2115/Ryan/r0058_l10/redshift0058"
pf = load(fn, file_style="%s.grid.cpu%%04i") # load data

val = 2.37163264206e-23
pos = [0.40328598,0.47176743,0.46131516]

#fields = ["HI_Number_Density","HAlpha_Emissivity","HAlpha_Emissivity_R","HAlpha_Voort_R"]
fields = ["HAlpha_Emissivity_R","HAlpha_Voort_R"]

proj = pf.h.proj(0,fields)
width = 200./pf['kpc']
res = [500,500]
frb = proj.to_frb(width,res,center=pos)

xL = np.arange(-250,250)*0.4
xL, yL = np.meshgrid(xL,xL)

#rp_Arc = radial_data(frb['HAlpha_Emissivity'],x=xL,y=yL)
rp_Ral = radial_data(frb['HAlpha_Emissivity_R'],x=xL,y=yL)
rp_Vor = radial_data(frb['HAlpha_Voort_R'],x=xL,y=yL)
#rp_hdens = radial_data(frb['HI_Number_Density'],x=xL,y=yL)

#thisindex = (rp_Arc.r/250. >0.01)
#plt.plot((rp_Arc.r/250.)[thisindex],rp_Arc.mean[thisindex])
#plt.yscale('log')
#plt.xscale('log')
#plt.xlabel('Radius / Rvir')
#plt.ylabel('Surface Brightness (ergs s^-1 cm^-2 arcsec^-2)')
#plt.savefig('rp_haArc_cloudy_smooth.png')
#plt.close()

thisindex = (rp_Ral.r/250. >0.01)
plt.plot((rp_Ral.r/250.)[thisindex],rp_Ral.mean[thisindex])
plt.yscale('log')
plt.xscale('log')
plt.xlabel('Radius / Rvir')
plt.ylabel('Surface Brightness (R)')
plt.xlim(10**-2,1)
plt.ylim(10**-5,10)
plt.savefig('rp_haR_cloudy_smooth.png')
plt.close()

thisindex = (rp_Vor.r/250. >0.01)
plt.plot((rp_Vor.r/250.)[thisindex],rp_Vor.mean[thisindex])
plt.yscale('log')
plt.xscale('log')
plt.xlabel('Radius / Rvir')
plt.ylabel('Surface Brightness (R)')
plt.xlim(10**-2,1)
plt.ylim(10**-5,10)
plt.savefig('rp_haVoortR_cloudy_smooth.png')
plt.close()

#thisindex = (rp_hdens.r/250. >0.01)
#plt.plot((rp_hdens.r/250.)[thisindex],rp_hdens.mean[thisindex])
#plt.yscale('log')
#plt.xscale('log')
#plt.xlabel('Radius / Rvir')
#plt.ylabel('HI Column Density (cm^-2)')
#plt.savefig('rp_hdens_cloudy_smooth.png')
#plt.close()


# for colab
# from google.colab import auth
from private_constants import PrivateConstants

import os
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload
from google.cloud import storage

class GCS:
    def list_files_from_gcs(bucket_name, path):
        gcs_service = storage.Client(PrivateConstants.PROJECT_NAME)
        bucket = gcs_service.bucket(bucket_name)
        file_names = []
        blobs = bucket.list_blobs(prefix=path)
        for blob in blobs:
            file_names.append(blob.name)
        return file_names

    def download_files_from_gcs(bucket_name, path, to_dir):
        gcs_service = storage.Client(PrivateConstants.PROJECT_NAME)
        bucket = gcs_service.bucket(bucket_name)
        blobs = bucket.list_blobs(prefix=path)
        for blob in blobs:
            file_name = blob.name.rsplit('/', 1)[-1]
            if file_name == "":
                continue
            blob.download_to_filename(os.path.join(to_dir, file_name))
        print("Download to {}/{} complete".format(bucket_name, path))
        return True

    def upload_to_gcs(bucket_name, from_file, to_file):
        gcs_service = storage.Client(PrivateConstants.PROJECT_NAME)
        bucket = gcs_service.bucket(bucket_name)
        blob = bucket.blob(to_file)
        blob.upload_from_filename(from_file)
        print("Upload to {}/{} complete".format(bucket_name, to_file))
        return True

    def download_file_from_gcs(bucket_name, from_file, to_file):
        gcs_service = storage.Client(PrivateConstants.PROJECT_NAME)
        bucket = gcs_service.bucket(bucket_name)
        blob = bucket.blob(from_file)
        blob.download_to_filename(to_file)
        print("Download from {}/{} complete".format(bucket_name, from_file))
        return True

    # DEPRECATED
    def get_from_gcs(bucket_name, from_file, to_file):
        # Create the service client.
        gcs_service = build('storage', 'v1')

        with open(to_file, 'wb') as f:
            request = gcs_service.objects().get_media(bucket=bucket_name,
                                                      object=from_file)
            media = MediaIoBaseDownload(f, request)

            done = False
            while not done:
                # _ is a placeholder for a progress object that we ignore.
                # (Our file is small, so we skip reporting progress.)
                _, done = media.next_chunk()

        print('Download complete')

    # DEPRECATED
    def store_in_gcs(bucket_name, from_file, to_file):
        # Create the service client.
        gcs_service = build('storage', 'v1')

        #   name for name in os.listdir(DIR) if os.path.isfile(os.path.join(DIR, name))

        with open(to_file, 'wb') as f:
            request = gcs_service.objects().get_media(bucket=bucket_name,
                                                      object=from_file)
            media = MediaIoBaseDownload(f, request)

            done = False
            while not done:
                # _ is a placeholder for a progress object that we ignore.
                # (Our file is small, so we skip reporting progress.)
                _, done = media.next_chunk()

        print('Upload complete')


import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os
from sklearn.preprocessing import MinMaxScaler
import time
plt.style.use("ggplot")
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
from keras.models import Sequential
from keras.layers import LSTM, Dense, Dropout
import joblib

start_time = time.time()
df = pd.read_csv('../../../../Crypto_Data/Data/bitcoin_new.csv', sep=',')
df = df.set_index("Date")[['Price']].tail(2500)
df = df.set_index(pd.to_datetime(df.index))

scaler = MinMaxScaler()
df = pd.DataFrame(scaler.fit_transform(df), columns=df.columns, index=df.index)

def visualize_training_results(results):
    """
    Plots the loss and accuracy for the training and testing data
    """
    history = results.history
    plt.figure(figsize=(12, 4))
    plt.plot(history['val_loss'])
    plt.plot(history['loss'])
    plt.legend(['val_loss', 'loss'])
    plt.title('Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.show()

    plt.figure(figsize=(12, 4))
    plt.plot(history['val_accuracy'])
    plt.plot(history['accuracy'])
    plt.legend(['val_accuracy', 'accuracy'])
    plt.title('Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.show()

def split_sequence(seq, n_steps_in, n_steps_out):
    x, y = [], []
    for i in range(len(seq)):
        end = i + n_steps_in
        out_end = end + n_steps_out
        if out_end > len(seq):
            break

        seq_x, seq_y = seq[i:end], seq[end:out_end]
        x.append(seq_x)
        y.append(seq_y)

    return np.array(x), np.array(y)

def layer_maker(n_layers, n_nodes, activation, drop=None, d_rate=.5):
    for x in range(1, n_layers + 1):
        model.add(LSTM(n_nodes, activation=activation, return_sequences=True))
        try:
            if x % drop == 0:
                model.add(Dropout(d_rate))
        except:
            pass

# Data prep
n_per_in = 30
n_per_out = 10
n_features = 1
x, y = split_sequence(list(df.Price), n_per_in, n_per_out)
x = x.reshape((x.shape[0], x.shape[1], n_features))

# Model
model = Sequential()
activ = "softsign"
# activ = "tanh"
# activ = "relu"
model.add(LSTM(30, activation=activ, return_sequences=True, input_shape=(n_per_in, n_features)))
layer_maker(n_layers=6, n_nodes=10, activation=activ)
model.add(LSTM(10, activation=activ))
model.add(Dense(n_per_out))
# model.summary()

model.compile(optimizer='adam', loss='mse', metrics=['accuracy'])
res = model.fit(x, y, epochs=100, batch_size=128, validation_split=0.1, verbose=0)
yhat = model.predict(x[-2].reshape(1, n_per_in, n_features)).tolist()[0]
yhat = scaler.inverse_transform(np.array(yhat).reshape(-1, 1)).tolist()
yhat = pd.DataFrame(yhat, index=pd.date_range(start=df.index[-n_per_out], periods=len(yhat), freq="D"), columns=df.columns)
# joblib.dump(model, './Models/rnn_9.sav')
# joblib.dump(yhat, './Models/yhat_9.sav')
actual = scaler.inverse_transform(y[-1].reshape(-1, 1))
# joblib.dump(actual, './Models/actual_9.sav')
yhat_fut = model.predict(np.array(df.tail(n_per_in)).reshape(1, n_per_in, n_features)).tolist()[0]
yhat_fut = scaler.inverse_transform(np.array(yhat_fut).reshape(-1,1)).tolist()
preds = pd.DataFrame(yhat_fut, index=pd.date_range(start=df.index[-1], periods=len(yhat_fut), freq="D"), columns=df.columns)
# joblib.dump(preds, './Models/preds_9.sav')
pers = 30
actual2 = pd.DataFrame(scaler.inverse_transform(df[["Price"]].tail(pers)), index=df.Price.tail(pers).index, columns=df.columns).append(preds.head(1))
# joblib.dump(actual2, './Models/actual2_8.sav')

print("--- %s  ---" % (time.time() - start_time))
visualize_training_results(res)


import pathlib
import sys
from textwrap import dedent
import shutil
import cProfile

from .errors import *

SERIAL = "serial"
MPI_PARALLEL = "mpi"

IN_PROGRESS_PREFIX = "inprogress_"

class PipelineStage:
    """A PipelineStage implements a single calculation step within a wider pipeline.

    Each different type of analysis stge is represented by a subclass of this
    base class.  The base class handles the connection between different pipeline
    stages, and the execution of the stages within a workflow system (parsl),
    potentially in parallel (MPI).

    An instance of one of these classes represents an actual run of the stage,
    with the required inputs, outputs, and configuration specified.

    See documentation pages for more details.

    """

    parallel = True
    config_options = {}
    doc = ""

    def __init__(self, args, comm=None):
        """Construct a pipeline stage, specifying the inputs, outputs, and configuration for it.

        The constructor needs a dict or namespace. It should include:
        - input paths (required)
        - config path (required)
        - output paths (optional but usual)
        - additional configuration (required if not specified elsewhere)

        Input and output paths should map tags to paths.
        Tags are strings, and the first elements in each item in the subclass's
        "inputs" and "output" attributes.
        e.g. for a subclass with:
            inputs = [('eggs', TextFile)]
            outputs = [('spam', TextFile)]
        the args could contain:
            {'eggs': 'inputs/eggs.txt',
             'spam': 'outputs/spam.txt' }
        If spam is not specified it will default to "./spam.txt"

        }

        The config should map "config" to a path where a YAML config file
        is located, e.g. {'config':'/path/to/config.yml'}

        Any config variables that are specified in the class's config attribute
        will be searched for first in args, then in the config file, and then
        by looking at any default value they have been given.
        If they have no default value (and just a type, like int, is listed), then
        it's an error if they are not specified somewhere.

        The execute method can instantiate and run the class together, with added bonuses
        like profiling and debugging tools.

        Parameters
        ----------
        args: dict or namespace
            Specification of input and output paths and any missing config options
        comm: MPI communicator
            (default is None) An MPI comm object to use in preference to COMM_WORLD
        """
        if not isinstance(args, dict):
            args = vars(args)

        # We first check for missing input files, that's a show stopper
        missing_inputs = []
        for x in self.input_tags():
            val = args[x]
            if val is None:
                missing_inputs.append(f"--{x}")
        if missing_inputs:
            missing_inputs = "  ".join(missing_inputs)
            raise ValueError(
                f"""

Missing these names on the command line:
    Input names: {missing_inputs}"""
            )

        self._inputs = {x: args[x] for x in self.input_tags()}
        # We alwys assume the config arg exists, whether it is in input_tags or not
        if args.get("config") is None:
            raise ValueError("The argument --config was missing on the command line.")

        self._inputs["config"] = args["config"]

        # We prefer to receive explicit filenames for the outputs but will
        # tolerate missing output filenames and will default to tag name in
        # current folder (this is for CWL compliance)
        self._outputs = {}
        for i, x in enumerate(self.output_tags()):
            if args[x] is None:
                ftype = self.outputs[i][1]
                self._outputs[x] = ftype.make_name(x)
            else:
                self._outputs[x] = args[x]

        # Finally, we extract configuration information from a combination of
        # command line arguments and optional 'config' file
        self._configs = self.read_config(args)

        use_mpi = args.get("mpi", False)
        if use_mpi:
            try:
                # This isn't a ceci dependency, so give a sensible error message if not installed.
                import mpi4py.MPI
            except ImportError:
                print("ERROR: Using --mpi option requires mpi4py to be installed.")
                raise

        # For scripting and testing we allow an MPI communicator or anything
        # with the same API to be passed in directly, overriding the --mpi
        # flag.
        if comm is not None:
            self._parallel = MPI_PARALLEL
            self._comm = comm
            self._size = self._comm.Get_size()
            self._rank = self._comm.Get_rank()
        elif use_mpi:
            self._parallel = MPI_PARALLEL
            self._comm = mpi4py.MPI.COMM_WORLD
            self._size = self._comm.Get_size()
            self._rank = self._comm.Get_rank()
        else:
            self._parallel = SERIAL
            self._comm = None
            self._size = 1
            self._rank = 0

    pipeline_stages = {}
    incomplete_pipeline_stages = {}

    def __init_subclass__(cls, **kwargs):
        """
        Python 3.6+ provides a facility to automatically
        call a method (this one) whenever a new subclass
        is defined.  In this case we use that feature to keep
        track of all available pipeline stages, each of which is
        defined by a class.

        """
        super().__init_subclass__(**kwargs)

        # This is a hacky way of finding the file
        # where our stage was defined
        filename = sys.modules[cls.__module__].__file__

        stage_is_complete = (
            hasattr(cls, "outputs") and hasattr(cls, "inputs") and hasattr(cls, "run")
        )

        # If there isn't an explicit name already then set it here.
        # by default use the class name.
        if not hasattr(cls, "name"):
            cls.name = cls.__name__

        if stage_is_complete:
            # Deal with duplicated class names
            if cls.name in cls.pipeline_stages:
                other = cls.pipeline_stages[cls.name][1]
                raise DuplicateStageName(
                    "You created two pipeline stages with the"
                    f"name {cls.name}.\nOne was in {filename}\nand the "
                    f"other in {other}\nYou can either change the class "
                    "name or explicitly put a variable 'name' in the top"
                    "level of the class."
                )

            # Check for "config" in the inputs list - this is implicit
            for name, _ in cls.inputs:
                if name == "config":
                    raise ReservedNameError(
                        "An input called 'config' is implicit in each pipeline "
                        "stage and should not be added explicitly.  Please update "
                        f"your pipeline stage called {cls.name} to remove/rename "
                        "the input called 'config'."
                    )

        # Check if user has over-written the config variable.
        # Quite a common error I make myself.
        if not isinstance(cls.config, property):
            raise ReservedNameError(
                "You have a class variable called 'config', which "
                "is reserved in ceci for its own configuration. "
                "You may have meant to specify config_options?"
            )
        # Find the absolute path to the class defining the file
        path = pathlib.Path(filename).resolve()

        # Register the class
        if stage_is_complete:
            cls.pipeline_stages[cls.name] = (cls, path)
        else:
            cls.incomplete_pipeline_stages[cls.__name__] = (cls, path)

    #############################################
    # Life cycle-related methods and properties.
    #############################################

    @classmethod
    def get_stage(cls, name):
        """
        Return the PipelineStage subclass with the given name.

        This is used so that we do not need a new entry point __main__ function
        for each new stage - instead we can just use a single one which can query
        which class it should be using based on the name.

        Returns
        -------
        cls: class
            The corresponding subclass
        """
        stage = cls.pipeline_stages.get(name)

        # If not found, then check for incomplete stages
        if stage is None:
            if name in cls.incomplete_pipeline_stages:
                raise IncompleteStage(
                    f"The stage {name} is not completely written. "
                    "Stages must specify 'inputs', 'outputs' as class variables "
                    f"and a 'run' method.\n{name} might be unfinished, or it might "
                    "be intended as a base for other classes and not to be run."
                )
            else:
                raise StageNotFound(f"Unknown stage '{name}'")
        else:
            return stage[0]

    @classmethod
    def get_module(cls):
        """
        Return the path to the python package containing the current sub-class

        If we have a PipelineStage subclass defined in a module called "bar", in
        a package called "foo" e.g.:
        /path/to/foo/bar.py  <--   contains subclass "Baz"

        Then calling Baz.get_module() will return "foo.bar".

        We use this later to construct command lines like "python -m foo Baz"

        Returns
        -------
        module: str
            The module containing this class.
        """
        return cls.pipeline_stages[cls.name][0].__module__

    @classmethod
    def usage(cls):
        """
        Print a usage message.
        """
        stage_names = "\n- ".join(cls.pipeline_stages.keys())
        try:
            module = cls.get_module().split(".")[0]
        except:
            module = "<module_name>"
        sys.stderr.write(
            f"""
Usage: python -m {module} <stage_name> <stage_arguments>

If no stage_arguments are given then usage information
for the chosen stage will be given.

I currently know about these stages:
- {stage_names}
"""
        )

    @classmethod
    def main(cls):
        """
        Create an instance of this stage and execute it with
        inputs and outputs taken from the command line
        """
        try:
            stage_name = sys.argv[1]
        except IndexError:
            cls.usage()
            return 1
        if stage_name in ["--help", "-h"] and len(sys.argv) == 2:
            cls.usage()
            return 1
        stage = cls.get_stage(stage_name)
        args = stage._parse_command_line()
        stage.execute(args)
        return 0

    @classmethod
    def _parse_command_line(cls):
        cmd = " ".join(sys.argv[:])
        import argparse

        parser = argparse.ArgumentParser(description=f"Run pipeline stage {cls.name}")
        parser.add_argument("stage_name")
        for conf in cls.config_options:
            def_val = cls.config_options[conf]
            opt_type = def_val if type(def_val) == type else type(def_val)

            if opt_type == bool:
                parser.add_argument(f"--{conf}", action="store_const", const=True)
            elif opt_type == list:
                out_type = def_val[0] if type(def_val[0]) == type else type(def_val[0])
                if out_type is str:
                    parser.add_argument(
                        f"--{conf}", type=lambda string: string.split(",")
                    )
                elif out_type is int:
                    parser.add_argument(
                        f"--{conf}",
                        type=lambda string: [int(i) for i in string.split(",")],
                    )
                elif out_type is float:
                    parser.add_argument(
                        f"--{conf}",
                        type=lambda string: [float(i) for i in string.split(",")],
                    )
                else:
                    raise NotImplementedError(
                        "Only handles str, int and float list arguments"
                    )
            else:
                parser.add_argument(f"--{conf}", type=opt_type)
        for inp in cls.input_tags():
            parser.add_argument(f"--{inp}")
        for out in cls.output_tags():
            parser.add_argument(f"--{out}")
        parser.add_argument("--config")

        if cls.parallel:
            parser.add_argument(
                "--mpi", action="store_true", help="Set up MPI parallelism"
            )
        parser.add_argument(
            "--pdb", action="store_true", help="Run under the python debugger"
        )
        parser.add_argument(
            "--cprofile",
            action="store",
            default="",
            type=str,
            help="Profile the stage using the python cProfile tool",
        )
        args = parser.parse_args()
        return args

    @classmethod
    def execute(cls, args):
        """
        Create an instance of this stage and run it
        with the specified inputs and outputs.

        This is calld by the main method.

        Parameters
        ----------
        args: namespace
            The argparse namespace for this subclass.
        """
        import pdb

        stage = cls(args)

        if stage.rank == 0:
            print(f"Executing stage: {cls.name}")

        if args.cprofile:
            profile = cProfile.Profile()
            profile.enable()

        try:
            stage.run()
        except Exception as error:
            if args.pdb:
                print(
                    "There was an exception - starting python debugger because you ran with --pdb"
                )
                print(error)
                pdb.post_mortem()
            else:
                raise
        try:
            stage.finalize()
        except Exception as error:
            if args.pdb:
                print(
                    "There was an exception in the finalization - starting python debugger because you ran with --pdb"
                )
                print(error)
                pdb.post_mortem()
            else:
                raise
        if args.cprofile:
            profile.disable()
            profile.dump_stats(args.cprofile)
            profile.print_stats("cumtime")

        if stage.rank == 0:
            print(f"Stage complete: {cls.name}")

    def finalize(self):
        """Finalize the stage, moving all its outputs to their final locations."""
        # Synchronize files so that everything is closed
        if self.is_mpi():
            self.comm.Barrier()

        # Move files to their final path
        # only the master process moves things
        if self.rank == 0:
            for tag in self.output_tags():
                # find the old and new names
                temp_name = self.get_output(tag)
                final_name = self.get_output(tag, final_name=True)

                # it's not an error here if the path does not exist,
                # because that will be handled later.
                if pathlib.Path(temp_name).exists():
                    # replace directories, rather than nesting more results
                    if pathlib.Path(final_name).is_dir():
                        shutil.rmtree(final_name)
                    shutil.move(temp_name, final_name)
                else:
                    sys.stderr.write(
                        f"NOTE/WARNING: Expected output file {final_name} was not generated.\n"
                    )

    #############################################
    # Parallelism-related methods and properties.
    #############################################

    @property
    def rank(self):
        """The rank of this process under MPI (0 if not running under MPI)"""
        return self._rank

    @property
    def size(self):
        """The number or processes under MPI (1 if not running under MPI)"""
        return self._size

    @property
    def comm(self):
        """The MPI communicator object (None if not running under MPI)"""
        return self._comm

    def is_parallel(self):
        """
        Returns True if the code is being run in parallel.
        Right now is_parallel() will return the same value as is_mpi(),
        but that may change in future if we implement other forms of
        parallelization.
        """
        return self._parallel != SERIAL

    def is_mpi(self):
        """
        Returns True if the stage is being run under MPI.
        """
        return self._parallel == MPI_PARALLEL

    def split_tasks_by_rank(self, tasks):
        """Iterate through a list of items, yielding ones this process is responsible for/

        Tasks are allocated in a round-robin way.

        Parameters
        ----------
        tasks: iterable
            Tasks to split up

        """
        for i, task in enumerate(tasks):
            if i % self.size == self.rank:
                yield task

    def data_ranges_by_rank(self, n_rows, chunk_rows, parallel=True):
        """Split a number of rows by process.

        Given a total number of rows to read and a chunk size, yield
        the ranges within them that this process should handle.

        Parameters
        ----------
        n_rows: int
            Total number of rows to split up

        chunk_rows: int
            Size of each chunk to be read.

        Parallel: bool
            Whether to split data by rank or just give all procs all data.
            Default=True
        """
        n_chunks = n_rows // chunk_rows
        if n_chunks * chunk_rows < n_rows:
            n_chunks += 1
        if parallel:
            it = self.split_tasks_by_rank(range(n_chunks))
        else:
            it = range(n_chunks)
        for i in it:
            start = i * chunk_rows
            end = min((i + 1) * chunk_rows, n_rows)
            yield start, end

    ##################################################
    # Input and output-related methods and properties.
    ##################################################

    def get_input(self, tag):
        """Return the path of an input file with the given tag"""
        return self._inputs[tag]

    def get_output(self, tag, final_name=False):
        """Return the path of an output file with the given tag

        If final_name is False then use a temporary name - file will
        be moved to its final name at the end
        """
        path = self._outputs[tag]

        # If not the final version, add a tag at the start of the filename
        if not final_name:
            p = pathlib.Path(path)
            p = p.parent / (IN_PROGRESS_PREFIX + p.name)
            path = str(p)
        return path

    def open_input(self, tag, wrapper=False, **kwargs):
        """
        Find and open an input file with the given tag, in read-only mode.

        For general files this will simply return a standard
        python file object.

        For specialized file types like FITS or HDF5 it will return
        a more specific object - see the types.py file for more info.

        """
        path = self.get_input(tag)
        input_class = self.get_input_type(tag)
        obj = input_class(path, "r", **kwargs)

        if wrapper:
            return obj
        else:
            return obj.file

    def open_output(self, tag, wrapper=False, final_name=False, **kwargs):
        """
        Find and open an output file with the given tag, in write mode.

        If final_name is True then they will be opened using their final
        target output name.  Otherwise we will prepend "inprogress_" to their
        file name. This means we know that if the final file exists then it
        is completed.

        If wrapper is True this will return an instance of the class
        of the file as specified in the cls.outputs.  Otherwise it will
        return an open file object (standard python one or something more
        specialized).

        Parameters
        ----------

        tag: str
            Tag as listed in self.outputs

        wrapper: bool
            Default=False.  Whether to return a wrapped file

        final_name: bool
            Default=False. Whether to save to

        **kwargs:
            Extra args are passed on to the file's class constructor.

        """
        path = self.get_output(tag, final_name=final_name)
        output_class = self.get_output_type(tag)

        # HDF files can be opened for parallel writing
        # under MPI.  This checks if:
        # - we have been told to open in parallel
        # - we are actually running under MPI
        # and adds the flags required if all these are true
        run_parallel = kwargs.pop("parallel", False) and self.is_mpi()
        if run_parallel:
            kwargs["driver"] = "mpio"
            kwargs["comm"] = self.comm

            # XXX: This is also not a dependency, but it should be.
            #      Or even better would be to make it a dependency of descformats where it
            #      is actually used.
            import h5py

            if not h5py.get_config().mpi:
                print(
                    dedent(
                        """\
                Your h5py installation is not MPI-enabled.
                Options include:
                  1) Set nprocess to 1 for all stages
                  2) Upgrade h5py to use mpi.  See instructions here:
                     http://docs.h5py.org/en/latest/build.html#custom-installation
                Note: If using conda, the most straightforward way is to enable it is
                    conda install -c spectraldns h5py-parallel
                """
                    )
                )
                raise RuntimeError("h5py module is not MPI-enabled.")

        # Return an opened object representing the file
        obj = output_class(path, "w", **kwargs)
        if wrapper:
            return obj
        else:
            return obj.file

    @classmethod
    def output_tags(cls):
        """
        Return the list of output tags required by this stage
        """
        return [tag for tag, _ in cls.outputs]

    @classmethod
    def input_tags(cls):
        """
        Return the list of input tags required by this stage
        """
        return [tag for tag, _ in cls.inputs]

    def get_input_type(self, tag):
        """Return the file type class of an input file with the given tag."""
        for t, dt in self.inputs:
            if t == tag:
                return dt
        raise ValueError(f"Tag {tag} is not a known input")

    def get_output_type(self, tag):
        """Return the file type class of an output file with the given tag."""
        for t, dt in self.outputs:
            if t == tag:
                return dt
        raise ValueError(f"Tag {tag} is not a known output")

    ##################################################
    # Configuration-related methods and properties.
    ##################################################

    @property
    def config(self):
        """
        Returns the configuration dictionary for this stage, aggregating command
        line options and optional configuration file.
        """
        return self._configs

    def read_config(self, args):
        """
        This function looks for the arguments of the pipeline stage using a
        combination of default values, command line options and separate
        configuration file.

        The order for resolving config options is first looking for a default
        value, then looking for a

        In case a mandatory argument (argument with no default) is missing,
        an exception is raised.

        Note that we recognize arguments with no default as the ones where
        self.config_options holds a type instead of a value.
        """
        # Try to load configuration file if provided
        import yaml

        # This is all the config information in the file, including
        # things for other stages
        overall_config = yaml.safe_load(open(self.get_input("config")))

        # The user can define global options that are inherited by
        # all the other sections if not already specified there.
        input_config = overall_config.get("global", {})

        # This is just the config info in the file for this stage.
        # It may be incomplete - there may be things specified on the
        # command line instead, or just using their default values
        stage_config = overall_config.get(self.name, {})
        input_config.update(stage_config)

        # Here we build up the actual configuration we use on this
        # run from all these sources
        my_config = {}

        # Loop over the options of the pipeline stage
        for x in self.config_options:
            opt = None
            opt_type = None

            # First look for a default value,
            # if a type (like int) is provided as the default it indicates that
            # this option doesn't have a default (i.e. is mandatory) and should
            # be explicitly provided with the specified type
            if type(self.config_options[x]) is type:
                opt_type = self.config_options[x]

            elif type(self.config_options[x]) is list:
                v = self.config_options[x][0]
                if type(v) is type:
                    opt_type = v
                else:
                    opt = self.config_options[x]
                    opt_type = type(v)
            else:
                opt = self.config_options[x]
                opt_type = type(opt)

            # Second, look for the option in the configuration file and override
            # default if provided TODO: Check types
            if x in input_config:
                opt = input_config[x]

            # Finally check for command line option that would override the value
            # in the configuration file. Note that the argument parser should
            # already have taken care of type
            if args[x] is not None:
                opt = args[x]

            # Finally, check that we got at least some value for this option
            if opt is None:
                raise ValueError(
                    f"Missing configuration option {x} for stage {self.name}"
                )

            my_config[x] = opt

        # Unspecified parameters can also be copied over.
        # This will be needed for parameters that are more complicated, such
        # as dictionaries or other more structured parameter information.
        for x, val in input_config.items():
            # Omit things we've already dealt with above
            if x in self.config_options:
                continue
            # copy over everything else
            else:
                my_config[x] = val

        return my_config

    def iterate_fits(self, tag, hdunum, cols, chunk_rows, parallel=True):
        """
        Loop through chunks of the input data from a FITS file with the given tag

        TODO: add ceci tests of this functions
        Parameters
        ----------
        tag: str
            The tag from the inputs list to use

        hdunum: int
            The extension number to read

        cols: list
            The columns to read

        chunk_rows: int
            Number of columns to read and return at once

        parallel: bool
            Whether to split up data among processes (parallel=True) or give
            all processes all data (parallel=False).  Default = True.

        Returns
        -------
        it: iterator
            Iterator yielding (int, int, array) tuples of (start, end, data)
            data is a structured array.
        """
        fits = self.open_input(tag)
        ext = fits[hdunum]
        n = ext.get_nrows()
        for start, end in self.data_ranges_by_rank(n, chunk_rows, parallel=True):
            data = ext.read_columns(cols, rows=range(start, end))
            yield start, end, data

    def iterate_hdf(self, tag, group_name, cols, chunk_rows, parallel=True):
        """
        Loop through chunks of the input data from an HDF5 file with the given tag.

        All the selected columns must have the same length.

        Parameters
        ----------
        tag: str
            The tag from the inputs list to use

        group: str
            The group within the HDF5 file to use, looked up as
            file[group]

        cols: list
            The columns to read

        chunk_rows: int
            Number of columns to read and return at once

        parallel: bool
            Whether to split up data among processes (parallel=True) or give
            all processes all data (parallel=False).  Default = True.

        Returns
        -------
        it: iterator
            Iterator yielding (int, int, dict) tuples of (start, end, data)
        """
        import numpy as np

        hdf = self.open_input(tag)
        group = hdf[group_name]

        # Check all the columns are the same length
        N = [len(group[col]) for col in cols]
        n = N[0]
        if not np.equal(N, n).all():
            raise ValueError(
                f"Different columns among {cols} in file {tag}\
            group {group_name} are different sizes - cannot use iterate_hdf"
            )

        # Iterate through the data providing chunks
        for start, end in self.data_ranges_by_rank(n, chunk_rows, parallel=parallel):
            data = {col: group[col][start:end] for col in cols}
            yield start, end, data

    ################################
    # Pipeline-related methods
    ################################

    @classmethod
    def generate_command(cls, inputs, config, outputs):
        """
        Generate a command line that will run the stage
        """
        module = cls.get_module()
        module = module.split(".")[0]

        flags = [cls.name]

        for tag, ftype in cls.inputs:
            try:
                fpath = inputs[tag]
            except KeyError:
                raise ValueError(f"Missing input location {tag}")
            flags.append(f"--{tag}={fpath}")

        flags.append(f"--config={config}")

        for tag, ftype in cls.outputs:
            try:
                fpath = outputs[tag]
            except:
                raise ValueError(f"Missing output location {tag}")
            flags.append(f"--{tag}={fpath}")

        flags = "   ".join(flags)

        # We just return this, instead of wrapping it in a
        # parsl job
        cmd = f"python3 -m {module} {flags}"
        return cmd

    @classmethod
    def generate_cwl(cls, log_dir=None):
        """
        Produces a CWL App object which can then be exported to yaml
        """
        import cwlgen

        module = cls.get_module()
        module = module.split(".")[0]

        # Basic definition of the tool
        cwl_tool = cwlgen.CommandLineTool(
            tool_id=cls.name,
            label=cls.name,
            base_command="python3",
            cwl_version="v1.0",
            doc=cls.__doc__,
        )
        if log_dir is not None:
            cwl_tool.stdout = f"{cls.name}.out"
            cwl_tool.stderr = f"{cls.name}.err"

        # Adds the first input binding with the name of the module and pipeline stage
        input_arg = cwlgen.CommandLineBinding(position=-1, value_from=f"-m{module}")
        cwl_tool.arguments.append(input_arg)
        input_arg = cwlgen.CommandLineBinding(position=0, value_from=f"{cls.name}")
        cwl_tool.arguments.append(input_arg)

        type_dict = {int: "int", float: "float", str: "string", bool: "boolean"}
        # Adds the parameters of the tool
        for opt in cls.config_options:
            def_val = cls.config_options[opt]

            # Handles special case of lists:
            if type(def_val) is list:
                v = def_val[0]
                param_type = {
                    "type": "array",
                    "items": type_dict[v] if type(v) == type else type_dict[type(v)],
                }
                default = def_val if type(v) != type else None
                input_binding = cwlgen.CommandLineBinding(
                    prefix="--{}=".format(opt), item_separator=",", separate=False
                )
            else:
                param_type = (
                    type_dict[def_val]
                    if type(def_val) == type
                    else type_dict[type(def_val)]
                )
                default = def_val if type(def_val) != type else None
                if param_type == "boolean":
                    input_binding = cwlgen.CommandLineBinding(prefix="--{}".format(opt))
                else:
                    input_binding = cwlgen.CommandLineBinding(
                        prefix="--{}=".format(opt), separate=False
                    )

            input_param = cwlgen.CommandInputParameter(
                opt,
                label=opt,
                param_type=param_type,
                input_binding=input_binding,
                default=default,
                doc="Some documentation about this parameter",
            )

            # We are bypassing the cwlgen builtin type check for the special case
            # of arrays until that gets added to the standard
            if type(def_val) is list:
                input_param.type = param_type

            cwl_tool.inputs.append(input_param)

        # Add the inputs of the tool
        for i, inp in enumerate(cls.input_tags()):
            input_binding = cwlgen.CommandLineBinding(prefix="--{}".format(inp))
            input_param = cwlgen.CommandInputParameter(
                inp,
                label=inp,
                param_type="File",
                param_format=cls.inputs[i][1].format,
                input_binding=input_binding,
                doc="Some documentation about the input",
            )
            cwl_tool.inputs.append(input_param)

        # Adds the overall configuration file
        input_binding = cwlgen.CommandLineBinding(prefix="--config")
        input_param = cwlgen.CommandInputParameter(
            "config",
            label="config",
            param_type="File",
            param_format="http://edamontology.org/format_3750",
            input_binding=input_binding,
            doc="Configuration file",
        )
        cwl_tool.inputs.append(input_param)

        # Add the definition of the outputs
        for i, out in enumerate(cls.output_tags()):
            output_name = cls.outputs[i][1].make_name(out)
            output_binding = cwlgen.CommandOutputBinding(glob=output_name)
            output = cwlgen.CommandOutputParameter(
                out,
                label=out,
                param_type="File",
                output_binding=output_binding,
                param_format=cls.outputs[i][1].format,
                doc="Some results produced by the pipeline element",
            )
            cwl_tool.outputs.append(output)

        if log_dir is not None:
            output = cwlgen.CommandOutputParameter(
                f"{cls.name}@stdout",
                label="stdout",
                param_type="stdout",
                doc="Pipeline elements standard output",
            )
            cwl_tool.outputs.append(output)
            error = cwlgen.CommandOutputParameter(
                f"{cls.name}@stderr",
                label="stderr",
                param_type="stderr",
                doc="Pipeline elements standard output",
            )
            cwl_tool.outputs.append(error)

        # Potentially add more metadata
        # This requires a schema however...
        # metadata = {'name': cls.name,
        #         'about': 'Some additional info',
        #         'publication': [{'id': 'one_doi'}, {'id': 'another_doi'}],
        #         'license': ['MIT']}
        # cwl_tool.metadata = cwlgen.Metadata(**metadata)

        return cwl_tool


"""
@author: Vincent Bonnet
@description : solve a 1D spring - aka damped harmonic oscillator
"""

import numpy as np
import matplotlib.pyplot as plt
from dataclasses import dataclass

SPRING_STIFFNESS = 1.0  # Newtons per meters
SPRING_DAMPING = 0.2   # Netwons per meters per second
INITIAL_POSITION = 2.0 # particle initial position in meters
MASS = 1.0  # particle mass in kg

# time constant
TIME_START = 0.0  # in Second
TIME_END = 30.0   # in Second
DT = 0.1         # in Second
N = round((TIME_END - TIME_START) / DT)

'''
 Classes (State, Derivate, Particle)
'''
@dataclass
class State:
    x = INITIAL_POSITION # position
    v = 0.0 # velocity

@dataclass
class Derivative:
    dx = 0.0 # derivative of x
    dv = 0.0 # derivative of v

@dataclass
class Particle:
    mass = 1.0
    state = State()

'''
  Helper Functions
'''
# internal spring force
def acceleration(state, mass):
    attachement = 0.0
    spring_force = -1.0 * (state.x - attachement) * SPRING_STIFFNESS  # spring force
    spring_force += -1.0 * (state.v * SPRING_DAMPING)  # spring damping
    acceleration = spring_force / mass
    return acceleration

# derivative of x - change in position
def dx(state):
    return state.v

# derivative of v - change in velocity
def dv(state, mass):
    return acceleration(state, mass)

# forward integration - return a new state
def integrate(state, derivative, dt):
    result_state = State()
    result_state.x = state.x + derivative.dx * dt
    result_state.v = state.v + derivative.dv * dt
    return result_state

# compute derivate at the state
def derivate(state, mass):
    result_derivate = Derivative()
    result_derivate.dx = dx(state)
    result_derivate.dv = dv(state, mass)
    return result_derivate

'''
 Integration Functions
'''
def forwardEuler(particle, time, dt):
    k = derivate(particle.state, particle.mass)
    particle.state = integrate(particle.state, k, dt)

def RK2(particle, time, dt):
    s1 = particle.state
    k1 = derivate(s1, particle.mass)
    s2 = integrate(s1, k1, dt * 0.5)
    k2 = derivate(s2, particle.mass)
    particle.state = integrate(particle.state, k2, dt)

def RK4(particle, time, dt):
    s1 = particle.state
    k1 = derivate(s1, particle.mass)
    s2 = integrate(s1, k1, dt * 0.5)
    k2 = derivate(s2, particle.mass)
    s3 = integrate(s1, k2, dt * 0.5)
    k3 = derivate(s3, particle.mass)
    s4 = integrate(s1, k3, dt)
    k4 = derivate(s4, particle.mass)
    k = Derivative()
    k.dx = k1.dx / 6 + k2.dx / 3 + k3.dx / 3 + k4.dx / 6
    k.dv = k1.dv / 6 + k2.dv / 3 + k3.dv / 3 + k4.dv / 6
    particle.state = integrate(particle.state, k, dt)

def semiImplicitEulerV1(particle, time, dt):
    particle.state.v += dv(particle.state, particle.mass) * dt
    particle.state.x += dx(particle.state) * dt

def semiImplicitEulerV2(particle, time, dt):
    particle.state.x += dx(particle.state) * dt
    particle.state.v += dv(particle.state, particle.mass) * dt

def leapFrog(particle, time, dt):
    if (time == TIME_START):
        # compute the velocity at half step which will cause the velocity to be half-step ahead of position
        particle.state.v += dv(particle.state, particle.mass) * dt * 0.5
    particle.state.x += particle.state.v * dt
    particle.state.v += dv(particle.state, particle.mass) * dt

def analyticSolution(particle, time, dt):
    w0 = np.sqrt(SPRING_STIFFNESS/particle.mass)
    y = SPRING_DAMPING / (2 * particle.mass)
    w = np.sqrt(w0 * w0 - y * y)
    a = np.exp(-1.0 * y * time)
    particle.state.x = (INITIAL_POSITION * a * np.cos(w * time));

def main():
    integrators = [(forwardEuler,"xkcd:aqua"),
                   (RK2,"xkcd:plum"),
                   (RK4,"xkcd:teal"),
                   (semiImplicitEulerV1,"xkcd:chartreuse"),
                   (semiImplicitEulerV2,"xkcd:olive"),
                   (leapFrog,"xkcd:green"),
                   (analyticSolution,"xkcd:red" )]

    plt.title('Single Damped Harmonic Oscillator')
    plt.xlabel('time(t)')
    plt.ylabel('position(x)')

    # integrators Loop
    time_samples = np.linspace(TIME_START, TIME_END, num=N, endpoint=False)
    position_samples = np.zeros(N)
    for integrator in integrators:
        function = integrator[0]
        plot_colour = integrator[1]
        particle = Particle();

        for i in range(N):
            position_samples[i] = particle.state.x
            function(particle, time_samples[i], DT)

        plt.plot(time_samples, position_samples, color=plot_colour, label=function.__name__)

    # display result
    plt.legend(bbox_to_anchor=(1, 1), loc=2)
    plt.show()

if __name__ == '__main__':
    main()