'''
This includes the module to realize the transformation between mesh and location map.
Part of the codes is adapted from https://github.com/Lotayou/densebody_pytorch
'''

from time import time
from tqdm import tqdm
from numpy.linalg import solve
from os.path import join
import torch
from torch import nn
import os
import torch.nn.functional as F

import pickle
import numpy as np
import pdb

'''
Index_UV_Generator is used to transform mesh and location map
The verts is in shape (B * V *C)
The UV map is in shape (B * H * W * C)
B: batch size;     V: vertex number;   C: channel number
H: height of uv map;  W: width of uv map
'''
class Index_UV_Generator(nn.Module):
    def __init__(self, UV_height, UV_width=-1, uv_type='BF', data_dir=None):
        super(Index_UV_Generator, self).__init__()

        if uv_type == 'SMPL':
            obj_file = 'smpl_fbx_template_expanded.obj'
        elif uv_type == 'BF':
            obj_file = 'smpl_boundry_free_template_expanded.obj'

        self.uv_type = uv_type

        if data_dir is None:
            d = os.path.dirname(__file__)
            data_dir = os.path.join(d, '..', 'data', 'uv_sampler')
        self.data_dir = data_dir
        self.h = UV_height
        self.w = self.h if UV_width < 0 else UV_width
        self.obj_file = obj_file
        self.para_file = 'paras_h{:04d}_w{:04d}_{}.npz'.format(self.h, self.w, self.uv_type)

        if not os.path.isfile(join(data_dir, self.para_file)):
            self.process()

        para = np.load(join(data_dir, self.para_file))

        self.v_index = torch.LongTensor(para['v_index'])
        self.bary_weights = torch.FloatTensor(para['bary_weights'])
        self.vt2v = torch.LongTensor(para['vt2v'])
        self.vt_count = torch.FloatTensor(para['vt_count'])
        self.texcoords = torch.FloatTensor(para['texcoords'])
        self.texcoords = 2 * self.texcoords - 1
        self.mask = torch.ByteTensor(para['mask'].astype('uint8'))

    def process(self):
        ##############################################################################
        # Load template obj file
        print('Loading obj file')
        with open(join(self.data_dir, self.obj_file), 'r') as fin:
            lines = [l
                     for l in fin.readlines()
                     if len(l.split()) > 0
                     and not l.startswith('#')
                     ]

        # Load all vertices (v) and texcoords (vt)
        vertices = []
        texcoords = []

        for line in lines:
            lsp = line.split()
            if lsp[0] == 'v':
                x = float(lsp[1])
                y = float(lsp[2])
                z = float(lsp[3])
                vertices.append((x, y, z))
            elif lsp[0] == 'vt':
                u = float(lsp[1])
                v = float(lsp[2])
                # texcoords.append((1 - v, u))
                texcoords.append((u, v))

        # Stack these into an array
        vertices = np.vstack(vertices).astype(np.float32)
        texcoords = np.vstack(texcoords).astype(np.float32)

        # Load face data. All lines are of the form:
        # f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3
        # Store the texcoord faces and a mapping from texcoord faces to vertex faces
        vt_faces = []
        v_num = vertices.shape[0]
        vt_num = texcoords.shape[0]
        vt2v = np.zeros(vt_num).astype('int64') - 1
        v2vt = [None] * v_num
        for i in range(v_num):
            v2vt[i] = set()

        for line in lines:
            vs = line.split()
            if vs[0] == 'f':
                v0 = int(vs[1].split('/')[0]) - 1
                v1 = int(vs[2].split('/')[0]) - 1
                v2 = int(vs[3].split('/')[0]) - 1
                vt0 = int(vs[1].split('/')[1]) - 1
                vt1 = int(vs[2].split('/')[1]) - 1
                vt2 = int(vs[3].split('/')[1]) - 1
                vt_faces.append((vt0, vt1, vt2))

                vt2v[vt0] = v0
                vt2v[vt1] = v1
                vt2v[vt2] = v2

                v2vt[v0].add(vt0)
                v2vt[v1].add(vt1)
                v2vt[v2].add(vt2)

        vt_faces = np.vstack(vt_faces)
        vt_count = np.zeros(v_num)
        for v_id in range(v_num):
            vt_count[v_id] = len(v2vt[v_id])

        ############################################################################
        # Calculating the barycentric weights used for UV map generation
        print('Calculating barycentric weights')
        s = time()
        h = self.h
        w = self.w
        face_num = vt_faces.shape[0]

        face_id = np.zeros((h, w), dtype=np.int)
        bary_weights = np.zeros((h, w, 3), dtype=np.float32)
        # uvs = texcoords * np.array([[h - 1, w - 1]])
        uvs = texcoords * np.array([[w - 1, h - 1]])
        grids = np.ones((face_num, 3), dtype=np.float32)
        anchors = np.concatenate((
            uvs[vt_faces].transpose(0, 2, 1),
            np.ones((face_num, 1, 3), dtype=uvs.dtype)
        ), axis=1)  # [F * 3 * 3]

        _loop = tqdm(np.arange(h * w), ncols=80)
        for i in _loop:
            r = i // w
            c = i % w
            # grids[:, 0] = r
            # grids[:, 1] = c

            grids[:, 0] = c
            grids[:, 1] = r

            weights = solve(anchors, grids)  # not enough accuracy?
            inside = np.logical_and.reduce(weights.T > 1e-10)
            index = np.where(inside == True)[0]

            if 0 == index.size:
                face_id[r, c] = -1  # just assign random id with all zero weights.
            # elif index.size > 1:
            #    print('bad %d' %i)
            else:
                face_id[r, c] = index[0]
                bary_weights[r, c] = weights[index[0]]

        v_index = vt2v[vt_faces[face_id]]
        mask = face_id >= 0
        v_index[face_id < 0, :] = 0
        bary_weights[face_id < 0, :] = 0
        print('Calculating finished. Time elapsed: {}s'.format(time() - s))

        ############################################################################
        # Ensure the neighboring pixels of vt on the UV map are meaningful
        tex = torch.FloatTensor(texcoords) * torch.FloatTensor([[w-1, h-1]])
        u_grid = tex[:, 0]
        u_lo = u_grid.floor().long().clamp(min=0, max=w - 1)
        u_hi = (u_lo + 1).clamp(max=w - 1)
        u_grid = torch.min(u_hi.float(), u_grid)
        u_w = u_grid - u_lo.float()

        v_grid = tex[:, 1]
        v_lo = v_grid.floor().long().clamp(min=0, max=h - 1)
        v_hi = (v_lo + 1).clamp(max=h - 1)
        v_grid = torch.min(v_hi.float(), v_grid)
        v_w = v_grid - v_lo.float()

        w_vlo_ulo = (1.0 - u_w) * (1.0 - v_w)
        w_vlo_uhi = u_w * (1.0 - v_w)
        w_vhi_ulo = (1.0 - u_w) * v_w
        w_vhi_uhi = u_w * v_w

        w = torch.cat([w_vlo_ulo, w_vlo_uhi, w_vhi_ulo, w_vhi_uhi], dim=0)
        u = torch.cat([u_lo, u_hi, u_lo, u_hi], dim=0)
        v = torch.cat([v_lo, v_lo, v_hi, v_hi], dim=0)
        v_id = torch.LongTensor(vt2v).repeat(4)        # The function repeat for ndarray and tensor are different!

        w_sorted, sort_index = w.sort(dim=0, descending=True)
        u_sorted = u[sort_index]
        v_sorted = v[sort_index]
        v_id_sorted = v_id[sort_index]

        # asign the empty pixel with the value of the 1 vt with maximal weights
        n_expand = 0
        for i in range(len(w_sorted)):
            m = mask[v_sorted[i], u_sorted[i]]
            if not m:
                v_index[v_sorted[i], u_sorted[i], 0] = v_id_sorted[i]
                bary_weights[v_sorted[i], u_sorted[i], 0] = 1
                mask[v_sorted[i], u_sorted[i]] = 1
                n_expand += 1

        np.savez(join(self.data_dir, self.para_file),
                 v_index=v_index,
                 bary_weights=bary_weights,
                 texcoords=texcoords,
                 vt2v=vt2v,
                 vt_count=vt_count,
                 mask=mask,
                 )

    def get_UV_map(self, verts):
        '''every pixel in UV map is determined by 3 adjacent vertices
        '''
        self.bary_weights = self.bary_weights.type(verts.dtype).to(verts.device) # shape: [res, res, 3]
        self.v_index = self.v_index.to(verts.device) # shape: [res, res, 3]

        if verts.dim() == 2:
            verts = verts.unsqueeze(0)

        im = verts[:, self.v_index, :]
        bw = self.bary_weights[:, :, None, :]
        im = torch.matmul(bw, im).squeeze(dim=3)
        return im

    def resample(self, uv_map):
        batch_size, _, _, channel_num = uv_map.shape
        v_num = self.vt_count.shape[0]
        self.texcoords = self.texcoords.type(uv_map.dtype).to(uv_map.device)
        self.vt2v = self.vt2v.to(uv_map.device)
        self.vt_count = self.vt_count.type(uv_map.dtype).to(uv_map.device)

        uv_grid = self.texcoords[None, None, :, :].expand(batch_size, -1, -1, -1)

        vt = F.grid_sample(uv_map.permute(0, 3, 1, 2), uv_grid, mode='bilinear')
        vt = vt.squeeze(2).permute(0, 2, 1)
        v = vt.new_zeros([batch_size, v_num, channel_num])
        v.index_add_(1, self.vt2v, vt)
        v = v / self.vt_count[None, :, None]
        return v

    # just used for the generation of GT UVmaps
    def forward(self, verts):
        return self.get_UV_map(verts)

    # generate part map
    def gen_part_map(self, face_part):
        part_map = np.zeros(self.bary_id.shape) - 1
        mask = self.bary_id >= 0
        part_map[mask] = face_part[self.bary_id[mask]]
        # plt.imshow(part_map)
        # part_map = self._dilate(torch.tensor(part_map+1)[None,:,:,None].float()).squeeze().numpy() -1
        return part_map

# Compute the weight map in UV space according to human body parts.
def cal_uv_weight(sampler, out_path):

    with open('data/segm_per_v_overlap.pkl', 'rb') as f:
        tmp = pickle.load(f)

    part_names = ['hips',
                  'leftUpLeg',
                  'rightUpLeg',
                  'spine',
                  'leftLeg',
                  'rightLeg',
                  'spine1',
                  'leftFoot',
                  'rightFoot',
                  'spine2',
                  'leftToeBase',
                  'rightToeBase',
                  'neck',
                  'leftShoulder',
                  'rightShoulder',
                  'head',
                  'leftArm',
                  'rightArm',
                  'leftForeArm',
                  'rightForeArm',
                  'leftHand',
                  'rightHand',
                  'leftHandIndex1',
                  'rightHandIndex1']

    part_weight = torch.tensor([1, 5, 5, 1, 5, 5, 1, 25, 25, 1, 25, 25, 2, 1, 1, 2, 5, 5, 5, 5, 25, 25, 25, 25])

    vert_part = torch.zeros([6890, 24])
    for i in range(24):
        key = part_names[i]
        verts = tmp[key]
        vert_part[verts, i] = 1

    part_map = sampler.get_UV_map(vert_part).squeeze(0)
    part_map = part_map > 0
    weight_map = part_weight[None, None, :].float() * part_map.float()
    weight_map = weight_map.max(dim=-1)[0]
    weight_map = weight_map / weight_map.mean()

    np.save(out_path, weight_map.numpy())
    return


#!/usr/bin/env python
"""
Program to generate sample SANS datasets for neural network training. A
modified version of compare_multi.py.

The program writes datafiles as result_<modelname>_<run_number> to the out/
directory. See example_data.dat for a sample of the file format used.

**TODOS**

* don't save q, dq with every dataset
* use hdf5 rather than sqlite

"""
from __future__ import print_function

from copy import deepcopy
import argparse
import os
import resource
import sys
import time
import traceback
from collections import OrderedDict, namedtuple
import sqlite3
import logging
import fnmatch

import numpy as np  # type: ignore

from sasmodels import core as sascore
from sasmodels import compare as sascomp
from sasmodels import data as sasdata

from . import sas_io
from .util.utils import columnize

MODELS = sascore.list_models()

# Maxim @ stackoverflow
# https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse/43357954#43357954
# (with mods by PAK)
def str2bool(value):
    """parse boolean argument"""
    if isinstance(value, bool):
       return value
    value = value.lower()
    if value in ('yes', 'true', 't', 'y', '1'):
        return True
    elif value in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value (1/0, Y[es]/N[o], T[rue]/F[alse]) expected.')

parser = argparse.ArgumentParser(
    description="""
    A script that generates SANS datasets for use in neural network training.""")
parser.add_argument(
    "models", nargs='*',
    help=f"A list of models or kinds ({', '.join(sascore.KINDS)})."
         f" A pattern such as *sphere will select all sphere models.")
parser.add_argument(
    "-x", "--exclude", type=str, default="",
    help="Exclude specific models separated by commas.")
parser.add_argument(
    "--tag", type=str, default="train",
    help="Tag for the generated data: train, test or validate.")
parser.add_argument(
    "--database", type=str, default=sas_io.DB_FILE,
    help="Path to the sqlite database file.")
parser.add_argument(
    "--count", type=int, default=1000,
    help="Count is the number of distinct models to generate.")
parser.add_argument(
    "--template", type=str, default="",
    help="SANS dataset defining q and resolution.")
parser.add_argument(
    "--resolution", type=float, default=3,
    help="Constant dQ/Q resolution percentage.")
parser.add_argument(
    "--noise", type=float, default=2,
    help="Constant dI/I uncertainty percentage.")
parser.add_argument(
    "--dimension", choices=('1D', '2D'), default='1D',
    help="Choose whether to generate 1D or 2D data.")
parser.add_argument(
    "--npoint", type=int, default=128,
    help="The number of points per model.")
parser.add_argument(
    "--mono", type=str2bool, default=True,
    help="Force all models to be monodisperse.")
parser.add_argument(
    "--magnetic", type=str2bool, default=False,
    help="Allow magnetic parameters in the model")
parser.add_argument(
    "--cutoff", type=float, default=0.,
    help="""
    CUTOFF is the cutoff value to use for the polydisperse distribution.
    Weights below the cutoff will be ignored.""")
parser.add_argument(
    "--precision", default='default',
    choices=['default', 'single', 'double', 'fast', 'single!', 'double!', 'quad!'],
    help="""
    Precision to use in floating point calculations. If postfixed with
    an '!', builds a DLL for the CPU. If default, use single unless the
    model requires double.""")
parser.add_argument(
    "-v", "--verbose",
    help="Verbose output level.", choices=[0, 1, 2])

# noinspection PyTypeChecker
def gen_data(model_name, data, count=1, noise=2,
             mono=True, magnetic=False, cutoff=1e-5,
             maxdim=np.inf, precision='double'):
    r"""
    Generates the data for the given model and parameters.

    *model_name* is the name of the model.

    *data* is the data object giving $q, \Delta q$ calculation points.

    *N* is the number of comparisons to make.

    *cutoff* is the polydispersity weight cutoff to make the calculation
    a little bit faster.

    *maxdim* is maximum value for any shape dimension.

    *precision* is the name of the calculation engine to use.

    Returns iterator *(seed, pars, data), ...* where *pars* is
    *{par: value, ...}* and *data* is *(q, dq, iq, diq)*.
    """
    is2d = False
    assert data.x.size > 0
    model_info = sascore.load_model_info(model_name)
    calculator = sascomp.make_engine(model_info, data, precision, cutoff)
    default_pars = sascomp.get_pars(model_info)
    assert calculator._data.x.size > 0
    x, dx = calculator._data.x, calculator._data.dx

    # A not very clean macro for evaluating the models, wich uses name and
    # seed from the current scope even though they haven't been defined yet.
    def simulate(pars):
        """
        Generate a random dataset for *fn* evaluated at *pars*.
        Returns *(x, dx, y, dy)*, o.

        Note that this replaces the data object within *fn*.
        """
        # TODO: support 2D data, which does not use x, dx, y, dy
        try:
            assert calculator._data.x.size > 0
            calculator.simulate_data(noise=noise, **pars)
            assert calculator._data.x.size > 0
            data = calculator._data
            # TODO: Do we need to copy? [Yes if data.y is reused.]
            result = (x, dx, data.y.copy(), data.dy.copy())
        except Exception:
            traceback.print_exc()
            print(f"Error when generating {model_name} for {seed}")
            result = (x, dx, np.NaN*x, np.NaN*x)
            #raise
        return result

    def pretty(pars):
        """
        Pretty the parameter set for displaying on one line
        """
        parlist = sascomp.parlist(model_info, pars, is2d)
        parlist = parlist.replace(os.linesep, '  ')
        parlist = parlist.replace(': ', '=')
        return parlist

    t0 = -np.inf
    interval = 5
    for k in range(count):
        seed = np.random.randint(int(1e6))
        t1 = time.perf_counter()
        if t1 > t0 + interval:
            print(f"generating {model_name} {k+1} of {count}")
            t0 = t1

        # Generate parameters
        with sascomp.push_seed(seed):
            pars = sascomp.randomize_pars(model_info, default_pars, maxdim)
        sascomp.constrain_pars(model_info, pars)
        if mono:
            pars = sascomp.suppress_pd(pars)
        if not magnetic:
            pars = sascomp.suppress_magnetism(pars)
        pars.update({'scale': 1, 'background': 1e-5})
        #print(f"{model_name} {seed} {pretty(pars)}")

        # Evaluate model
        data = simulate(pars) # q, dq, iq, diq

        # Skip data sets with NaN or negative numbers.
        # Note: some datasets will have fewer entries than others.
        if np.isnan(data[2]).any():
            print(f">>> NaN in {model_name} {seed} {pretty(pars)}")
            continue
        if (data[2] <= 0.).any():
            print(f">>> Negative values in {model_name} {seed} {pretty(pars)}")
            continue

        yield seed, pars, data

    # TODO: can free the calculator now
    print(f"Complete {model_name}")

def model_group(models, required=False):
    """
    Build a list of models from the items in *models*.  Could be individual
    model names or could be a unix-style glob pattern.
    """
    good = []
    bad = []
    for name in models:
        if name == "":
            continue
        items = fnmatch.filter(MODELS, name)
        if not items:
            try:
                items = sascore.list_models(name)
            except ValueError:
                pass
        if items:
            good.extend(items)
        else:
            bad.append(name)

    if bad:
        print(f"Bad model(s): {', '.join(bad)}.  ", file=sys.stderr)
    if bad or (required and not good):
        print(f"Use kind ({', '.join(sascore.KINDS)}) or one of:", file=sys.stderr)
        print(columnize(MODELS, indent="  "), file=sys.stderr, end='')
        print(f"Patterns such as *sphere will also work.")
        sys.exit(1)

    return sorted(set(good))

def run_model(opts):
    tag = opts.tag
    count = opts.count
    is2D = opts.dimension.startswith('2d')
    nq = opts.npoint
    mono = opts.mono
    magnetic = opts.magnetic
    cutoff = opts.cutoff if not mono else 0
    precision = opts.precision
    res = opts.resolution
    noise = opts.noise

    # Figure out which models we are using.
    include = model_group(opts.models, required=True)
    exclude = model_group(opts.exclude.split(','))
    model_list = sorted(set(include)-set(exclude))
    print("Selected models:\n", columnize(model_list, indent="  "))

    if opts.template:
        # Fetch q, dq from an actual SANS file
        data, index = sasdata.read(opts.template), None
    else:
        # Generate
        data, index = sascomp.make_data({
            'qmin': 1e-4, 'qmax': 0.2, 'is2d': is2D, 'nq': nq, 'res': res/100,
            'accuracy': 'Low', 'view': 'log', 'zero': False,
            })

    # Open database and lookup model counts.
    db = sas_io.sql_connect(opts.database)
    model_counts = sas_io.model_counts(db, tag)
    #print(model_counts)
    for model in model_list:
        # Figure out how many more entries we need for the model
        missing = count - model_counts.get(model, 0)
        if missing <= 0:
            continue
        # TODO: should not need deepcopy(data) but something is messing with q
        seq = gen_data(
            model, deepcopy(data), count=missing, mono=mono, magnetic=magnetic,
            cutoff=cutoff, precision=precision, noise=noise)
        # Process the missing entries batch by batch so if there is an
        # error we won't lose the entire group.
        for batch in chunk(seq, batch_size=100):
            sas_io.write_sql(db, model, batch, tag=tag)
            #sas_io.write_1d(path, model, items, tag=tag)
    #sas_io.read_sql(db, tag)
    db.close()

def chunk(seq, batch_size):
    """
    Generic iter tool: chunk sequence in groups of *batch_size*.

    Remaining items are returned in final group.
    """
    batch = []
    for item in seq:
        if len(batch) == batch_size:
            yield batch
            batch = []
        batch.append(item)
    yield batch

def main():
    logging.basicConfig(level=logging.INFO)
    opts = parser.parse_args()

    time_start = time.perf_counter()
    run_model(opts)
    time_end = time.perf_counter() - time_start
    print('Total computation time (s): %.2f' % (time_end / 10))
    print('Total memory usage: %.2f' %
          resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
    # Units of mem are OS dependent

if __name__ == "__main__":
    main()


"""
Honeycomb
=========

"""

import numpy as np

from .cof import Cof
from .vertices import LinearVertex, NonLinearVertex
from ..topology_graph import Edge

class Honeycomb(Cof):
    """
    Represents a honeycomb COF topology graph.

    Unoptimized construction

    .. moldoc::

        import moldoc.molecule as molecule
        import stk

        bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
        bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])

        cof = stk.ConstructedMolecule(
            topology_graph=stk.cof.Honeycomb(
                building_blocks=(bb1, bb2),
                lattice_size=(3, 3, 1),
            ),
        )
        moldoc_display_molecule = molecule.Molecule(
            atoms=(
                molecule.Atom(
                    atomic_number=atom.get_atomic_number(),
                    position=position,
                ) for atom, position in zip(
                    cof.get_atoms(),
                    cof.get_position_matrix(),
                )
            ),
            bonds=(
                molecule.Bond(
                    atom1_id=bond.get_atom1().get_id(),
                    atom2_id=bond.get_atom2().get_id(),
                    order=(
                        1
                        if bond.get_order() == 9
                        else bond.get_order()
                    ),
                ) for bond in cof.get_bonds()
            ),
        )

    ``Collapser(scale_steps=False)`` optimized construction

    .. moldoc::

        import moldoc.molecule as molecule
        import stk

        bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
        bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])

        cof = stk.ConstructedMolecule(
            topology_graph=stk.cof.Honeycomb(
                building_blocks=(bb1, bb2),
                lattice_size=(3, 3, 1),
                optimizer=stk.Collapser(scale_steps=False),
            ),
        )
        moldoc_display_molecule = molecule.Molecule(
            atoms=(
                molecule.Atom(
                    atomic_number=atom.get_atomic_number(),
                    position=position,
                ) for atom, position in zip(
                    cof.get_atoms(),
                    cof.get_position_matrix(),
                )
            ),
            bonds=(
                molecule.Bond(
                    atom1_id=bond.get_atom1().get_id(),
                    atom2_id=bond.get_atom2().get_id(),
                    order=(
                        1
                        if bond.get_order() == 9
                        else bond.get_order()
                    ),
                ) for bond in cof.get_bonds()
            ),
        )

    Building blocks with three and two functional groups are required
    for this topology graph.

    When using a :class:`dict` for the `building_blocks` parameter,
    as in :ref:`cof-topology-graph-examples`:
    *Multi-Building Block COF Construction*, a
    :class:`.BuildingBlock`, with the following number of functional
    groups, needs to be assigned to each of the following vertex ids:

        | 3-functional groups: 0 to 1
        | 2-functional groups: 2 to 4

    See :class:`.Cof` for more details and examples.

    """

    _lattice_constants = _a, _b, _c = (
        np.array([1., 0., 0.]),
        np.array([0.5, 0.866, 0]),
        np.array([0, 0, 5/1.7321])
    )

    _vertex_prototypes = (
        NonLinearVertex(0, (1/3)*_a + (1/3)*_b + (1/2)*_c),
        NonLinearVertex(1, (2/3)*_a + (2/3)*_b + (1/2)*_c),
    )

    _vertex_prototypes = (
        *_vertex_prototypes,
        LinearVertex.init_at_center(
            id=2,
            vertices=(_vertex_prototypes[0], _vertex_prototypes[1]),
        ),
        LinearVertex.init_at_shifted_center(
            id=3,
            vertices=(_vertex_prototypes[0], _vertex_prototypes[1]),
            cell_shifts=((0, 0, 0), (0, -1, 0)),
            lattice_constants=_lattice_constants,
        ),
        LinearVertex.init_at_shifted_center(
            id=4,
            vertices=(_vertex_prototypes[0], _vertex_prototypes[1]),
            cell_shifts=((0, 0, 0), (-1, 0, 0)),
            lattice_constants=_lattice_constants,
        )
    )

    _edge_prototypes = (
        Edge(0, _vertex_prototypes[2], _vertex_prototypes[0]),
        Edge(1, _vertex_prototypes[2], _vertex_prototypes[1]),

        Edge(2, _vertex_prototypes[3], _vertex_prototypes[0]),
        Edge(
            id=3,
            vertex1=_vertex_prototypes[3],
            vertex2=_vertex_prototypes[1],
            periodicity=(0, -1, 0),
        ),

        Edge(4, _vertex_prototypes[4], _vertex_prototypes[0]),
        Edge(
            id=5,
            vertex1=_vertex_prototypes[4],
            vertex2=_vertex_prototypes[1],
            periodicity=(-1, 0, 0),
        )
    )


#!/usr/bin/env python3
#
# Script to calculate the integral of one or multiple cube file
# by Patrick Melix
# 2022/04/04
#
# You can import the module and then call .main() or use it as a script
from curses import has_key
from pymatgen.io.vasp.outputs import Chgcar
from pymatgen.io.ase import AseAtomsAdaptor
from ase.io.cube import write_cube
import numpy as np
import os

def main(inFiles, outFiles, verbose=True, return_integrals=False, return_spin_integrals=False, mult_volume=False):
    assert len(inFiles) == len(outFiles), "Number of input and output files must be equal!"
    integrals = []
    spin_integrals = []
    for iFile,inFile in enumerate(inFiles):
        if not os.path.isfile(inFile):
            raise ValueError('File {:} does not exist'.format(inFile))

        #if output exists mv to .bak
        if os.path.isfile(outFiles[iFile]):
            if verbose: print('ATTENTION: {:} exists, moving to *.bak'.format(outFiles[iFile]))
            os.rename(outFiles[iFile], outFiles[iFile]+'.bak')

        if verbose: print("Reading {}".format(inFile))
        full_chgcar = Chgcar.from_file(inFile)
        spinpol = 'diff' in full_chgcar.data.keys()
        if return_spin_integrals and not spinpol:
            raise ValueError("File {} is not spinpolarized!".format(inFile))
        shape = full_chgcar.data['total'].shape
        n_data = np.prod(shape)

        if return_integrals:
            integrals.append(np.sum(np.abs(full_chgcar.data['total'])))
            integrals[-1] /= n_data
        if return_spin_integrals:
            spin_integrals.append(np.sum(np.abs(full_chgcar.data['diff'])))
            spin_integrals[-1] /= n_data
        if verbose:
            print("Shape of data: {}".format(shape))
            print("Total number of datapoints: {}".format(n_data))
            if return_integrals:
                integral = integrals[-1]
            else:
                integral = np.sum(np.abs(full_chgcar.data['total']))
                integral /= n_data
            print("Integral of total data is {}".format(integral))
            if spinpol:
                if return_spin_integrals:
                    spin_integral = spin_integrals[-1]
                else:
                    spin_integral = np.sum(np.abs(full_chgcar.data['diff']))
                    spin_integral /= n_data
                print("Integral of diff data is {}".format(spin_integral))

            origin = np.zeros(3)
            atoms = AseAtomsAdaptor.get_atoms(full_chgcar.structure)

            #Contrary to VASP Wiki, the CHGCAR is not rho*V, but rho*n_data.
            #So in order to have the integral over space = nelectrons, we need to divide by n_data.
            #Since this would result in super small numbers, we can transform to rho*V
            factor = n_data
            if mult_volume:
                factor /= atoms.get_volume()
            full_chgcar.data['total'] /= factor
            if spinpol: full_chgcar.data['diff'] /= factor
            #write cube
            filename = "{}.cube".format(outFiles[iFile])
            if verbose: print("Writing {}".format(filename))
            with open(filename, 'w') as f:
                write_cube(f, atoms, data=full_chgcar.data['total'], origin=origin)
            if spinpol:
                filename = "{}_mag.cube".format(outFiles[iFile])
                if verbose: print("Writing {}".format(filename))
                with open(filename, 'w') as f:
                    write_cube(f, atoms, data=full_chgcar.data['diff'], origin=origin)

    if return_integrals:
        if len(integrals) == 1:
            if return_spin_integrals: return integrals[0], spin_integrals[0]
            else: return integrals[0]
        else:
            if return_spin_integrals: return integrals, spinpol
            else: return integrals
    else:
        return

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Sum up cube files')
    parser.add_argument('input', type=str, nargs='+', help='Input Files')
    parser.add_argument('-output', type=str, nargs='+', help='Output File Names (no extension)')
    parser.add_argument('-v', help='Verbose', action='store_true')
    parser.add_argument('--volume', help='Multiply the Density with the Cell Volume', action='store_true')
    args = parser.parse_args()
    main(args.input, args.output, verbose=args.v, mult_volume=args.volume)


'''
Build a simple neural machine translation model
'''
import theano
import theano.tensor as tensor
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams

import cPickle as pkl
import ipdb
import numpy
import copy

import os
import warnings
import sys
import time

from collections import OrderedDict

from data_iterator import TextIterator

profile = False

# push parameters to Theano shared variables
def zipp(params, tparams):
    for kk, vv in params.iteritems():
        tparams[kk].set_value(vv)

# pull parameters from Theano shared variables
def unzip(zipped):
    new_params = OrderedDict()
    for kk, vv in zipped.iteritems():
        new_params[kk] = vv.get_value()
    return new_params

# get the list of parameters: Note that tparams must be OrderedDict
def itemlist(tparams):
    return [vv for kk, vv in tparams.iteritems()]

# dropout
def dropout_layer(state_before, use_noise, trng):
    proj = tensor.switch(
        use_noise,
        state_before * trng.binomial(state_before.shape, p=0.5, n=1,
                                     dtype=state_before.dtype),
        state_before * 0.5)
    return proj

# make prefix-appended name
def _p(pp, name):
    return '%s_%s' % (pp, name)

# initialize Theano shared variables according to the initial parameters
def init_tparams(params):
    tparams = OrderedDict()
    for kk, pp in params.iteritems():
        tparams[kk] = theano.shared(params[kk], name=kk)
    return tparams

# load parameters
def load_params(path, params):
    pp = numpy.load(path)
    for kk, vv in params.iteritems():
        if kk not in pp:
            warnings.warn('%s is not in the archive' % kk)
            continue
        params[kk] = pp[kk]

    return params

# layers: 'name': ('parameter initializer', 'feedforward')
layers = {'ff': ('param_init_fflayer', 'fflayer'),
          'gru': ('param_init_gru', 'gru_layer'),
          'gru_cond_simple': ('param_init_gru_cond_simple',
                              'gru_cond_simple_layer'),
          }

def get_layer(name):
    fns = layers[name]
    return (eval(fns[0]), eval(fns[1]))

# some utilities
def ortho_weight(ndim):
    W = numpy.random.randn(ndim, ndim)
    u, s, v = numpy.linalg.svd(W)
    return u.astype('float32')

def norm_weight(nin, nout=None, scale=0.01, ortho=True):
    if nout is None:
        nout = nin
    if nout == nin and ortho:
        W = ortho_weight(nin)
    else:
        W = scale * numpy.random.randn(nin, nout)
    return W.astype('float32')

def tanh(x):
    return tensor.tanh(x)

def linear(x):
    return x

def concatenate(tensor_list, axis=0):
    """
    Alternative implementation of `theano.tensor.concatenate`.
    This function does exactly the same thing, but contrary to Theano's own
    implementation, the gradient is implemented on the GPU.
    Backpropagating through `theano.tensor.concatenate` yields slowdowns
    because the inverse operation (splitting) needs to be done on the CPU.
    This implementation does not have that problem.
    :usage:
        >>> x, y = theano.tensor.matrices('x', 'y')
        >>> c = concatenate([x, y], axis=1)
    :parameters:
        - tensor_list : list
            list of Theano tensor expressions that should be concatenated.
        - axis : int
            the tensors will be joined along this axis.
    :returns:
        - out : tensor
            the concatenated tensor expression.
    """
    concat_size = sum(tt.shape[axis] for tt in tensor_list)

    output_shape = ()
    for k in range(axis):
        output_shape += (tensor_list[0].shape[k],)
    output_shape += (concat_size,)
    for k in range(axis + 1, tensor_list[0].ndim):
        output_shape += (tensor_list[0].shape[k],)

    out = tensor.zeros(output_shape)
    offset = 0
    for tt in tensor_list:
        indices = ()
        for k in range(axis):
            indices += (slice(None),)
        indices += (slice(offset, offset + tt.shape[axis]),)
        for k in range(axis + 1, tensor_list[0].ndim):
            indices += (slice(None),)

        out = tensor.set_subtensor(out[indices], tt)
        offset += tt.shape[axis]

    return out

# batch preparation, returns padded batches for both source and target
# sequences with their corresponding masks
def prepare_data(seqs_x, seqs_y, maxlen=None,
                 n_words_src=30000, n_words=30000):

    # x: a list of sentences
    lengths_x = [len(s) for s in seqs_x]
    lengths_y = [len(s) for s in seqs_y]

    # filter sequences according to maximum sequence length
    if maxlen is not None:
        new_seqs_x = []
        new_seqs_y = []
        new_lengths_x = []
        new_lengths_y = []
        for l_x, s_x, l_y, s_y in zip(lengths_x, seqs_x, lengths_y, seqs_y):
            if l_x < maxlen and l_y < maxlen:
                new_seqs_x.append(s_x)
                new_lengths_x.append(l_x)
                new_seqs_y.append(s_y)
                new_lengths_y.append(l_y)
        lengths_x = new_lengths_x
        seqs_x = new_seqs_x
        lengths_y = new_lengths_y
        seqs_y = new_seqs_y

        if len(lengths_x) < 1 or len(lengths_y) < 1:
            return None, None, None, None

    n_samples = len(seqs_x)
    maxlen_x = numpy.max(lengths_x) + 1
    maxlen_y = numpy.max(lengths_y) + 1

    # pad batches and create masks
    x = numpy.zeros((maxlen_x, n_samples)).astype('int64')
    y = numpy.zeros((maxlen_y, n_samples)).astype('int64')
    x_mask = numpy.zeros((maxlen_x, n_samples)).astype('float32')
    y_mask = numpy.zeros((maxlen_y, n_samples)).astype('float32')
    for idx, [s_x, s_y] in enumerate(zip(seqs_x, seqs_y)):
        x[:lengths_x[idx], idx] = s_x
        x_mask[:lengths_x[idx]+1, idx] = 1.
        y[:lengths_y[idx], idx] = s_y
        y_mask[:lengths_y[idx]+1, idx] = 1.

    return x, x_mask, y, y_mask

# feedforward layer: affine transformation + point-wise nonlinearity
def param_init_fflayer(options, params, prefix='ff', nin=None, nout=None,
                       ortho=True):
    if nin is None:
        nin = options['dim_proj']
    if nout is None:
        nout = options['dim_proj']
    params[_p(prefix, 'W')] = norm_weight(nin, nout, scale=0.01, ortho=ortho)
    params[_p(prefix, 'b')] = numpy.zeros((nout,)).astype('float32')

    return params

def fflayer(tparams, state_below, options, prefix='rconv',
            activ='lambda x: tensor.tanh(x)', **kwargs):
    return eval(activ)(
        tensor.dot(state_below, tparams[_p(prefix, 'W')]) +
        tparams[_p(prefix, 'b')])

# GRU layer
def param_init_gru(options, params, prefix='gru', nin=None, dim=None):

    if nin is None:
        nin = options['dim_proj']
    if dim is None:
        dim = options['dim_proj']

    # embedding to gates transformation weights, biases
    W = numpy.concatenate([norm_weight(nin, dim),
                           norm_weight(nin, dim)], axis=1)
    params[_p(prefix, 'W')] = W
    params[_p(prefix, 'b')] = numpy.zeros((2 * dim,)).astype('float32')

    # recurrent transformation weights for gates
    U = numpy.concatenate([ortho_weight(dim),
                           ortho_weight(dim)], axis=1)
    params[_p(prefix, 'U')] = U

    # embedding to hidden state proposal weights, biases
    Wx = norm_weight(nin, dim)
    params[_p(prefix, 'Wx')] = Wx
    params[_p(prefix, 'bx')] = numpy.zeros((dim,)).astype('float32')

    # recurrent transformation weights for hidden state proposal
    Ux = ortho_weight(dim)
    params[_p(prefix, 'Ux')] = Ux

    return params

def gru_layer(tparams, state_below, options, prefix='gru', mask=None,
              **kwargs):
    nsteps = state_below.shape[0]
    if state_below.ndim == 3:
        n_samples = state_below.shape[1]
    else:
        n_samples = 1

    dim = tparams[_p(prefix, 'Ux')].shape[1]

    if mask is None:
        mask = tensor.alloc(1., state_below.shape[0], 1)

    # utility function to slice a tensor
    def _slice(_x, n, dim):
        if _x.ndim == 3:
            return _x[:, :, n*dim:(n+1)*dim]
        return _x[:, n*dim:(n+1)*dim]

    # state_below is the input word embeddings
    # input to the gates, concatenated
    state_below_ = tensor.dot(state_below, tparams[_p(prefix, 'W')]) + \
        tparams[_p(prefix, 'b')]
    # input to compute the hidden state proposal
    state_belowx = tensor.dot(state_below, tparams[_p(prefix, 'Wx')]) + \
        tparams[_p(prefix, 'bx')]

    # step function to be used by scan
    # arguments    | sequences |outputs-info| non-seqs
    def _step_slice(m_, x_, xx_,  h_,          U, Ux):
        preact = tensor.dot(h_, U)
        preact += x_

        # reset and update gates
        r = tensor.nnet.sigmoid(_slice(preact, 0, dim))
        u = tensor.nnet.sigmoid(_slice(preact, 1, dim))

        # compute the hidden state proposal
        preactx = tensor.dot(h_, Ux)
        preactx = preactx * r
        preactx = preactx + xx_

        # hidden state proposal
        h = tensor.tanh(preactx)

        # leaky integrate and obtain next hidden state
        h = u * h_ + (1. - u) * h
        h = m_[:, None] * h + (1. - m_)[:, None] * h_

        return h

    # prepare scan arguments
    seqs = [mask, state_below_, state_belowx]
    init_states = [tensor.alloc(0., n_samples, dim)]
    _step = _step_slice
    shared_vars = [tparams[_p(prefix, 'U')],
                   tparams[_p(prefix, 'Ux')]]

    rval, updates = theano.scan(_step,
                                sequences=seqs,
                                outputs_info=init_states,
                                non_sequences=shared_vars,
                                name=_p(prefix, '_layers'),
                                n_steps=nsteps,
                                profile=profile,
                                strict=True)
    rval = [rval]
    return rval

# Conditional GRU layer without Attention
def param_init_gru_cond_simple(options, params, prefix='gru_cond', nin=None,
                               dim=None, dimctx=None):
    if nin is None:
        nin = options['dim']
    if dim is None:
        dim = options['dim']
    if dimctx is None:
        dimctx = options['dim']

    params = param_init_gru(options, params, prefix, nin=nin, dim=dim)

    # context to GRU gates
    Wc = norm_weight(dimctx, dim*2)
    params[_p(prefix, 'Wc')] = Wc

    # context to hidden proposal
    Wcx = norm_weight(dimctx, dim)
    params[_p(prefix, 'Wcx')] = Wcx

    return params

def gru_cond_simple_layer(tparams, state_below, options, prefix='gru',
                          mask=None, context=None, one_step=False,
                          init_state=None,
                          **kwargs):

    assert context, 'Context must be provided'

    if one_step:
        assert init_state, 'previous state must be provided'

    nsteps = state_below.shape[0]
    if state_below.ndim == 3:
        n_samples = state_below.shape[1]
    else:
        n_samples = 1

    # mask
    if mask is None:
        mask = tensor.alloc(1., state_below.shape[0], 1)

    dim = tparams[_p(prefix, 'Ux')].shape[1]

    # initial/previous state
    if init_state is None:
        init_state = tensor.alloc(0., n_samples, dim)

    assert context.ndim == 2, 'Context must be 2-d: #sample x dim'
    # projected context to GRU gates
    pctx_ = tensor.dot(context, tparams[_p(prefix, 'Wc')])
    # projected context to hidden state proposal
    pctxx_ = tensor.dot(context, tparams[_p(prefix, 'Wcx')])

    def _slice(_x, n, dim):
        if _x.ndim == 3:
            return _x[:, :, n*dim:(n+1)*dim]
        return _x[:, n*dim:(n+1)*dim]

    # projected x to gates
    state_belowx = tensor.dot(state_below, tparams[_p(prefix, 'Wx')]) + \
        tparams[_p(prefix, 'bx')]
    # projected x to hidden state proposal
    state_below_ = tensor.dot(state_below, tparams[_p(prefix, 'W')]) + \
        tparams[_p(prefix, 'b')]

    # step function to be used by scan
    # arguments    | sequences |outputs-info|       non-seqs
    def _step_slice(m_, x_, xx_,    h_,       pctx_, pctxx_, U, Ux):
        preact = tensor.dot(h_, U)
        preact += x_
        preact += pctx_
        preact = tensor.nnet.sigmoid(preact)

        r = _slice(preact, 0, dim)
        u = _slice(preact, 1, dim)

        preactx = tensor.dot(h_, Ux)
        preactx *= r
        preactx += xx_
        preactx += pctxx_

        h = tensor.tanh(preactx)

        h = u * h_ + (1. - u) * h
        h = m_[:, None] * h + (1. - m_)[:, None] * h_

        return h

    seqs = [mask, state_below_, state_belowx]
    _step = _step_slice

    shared_vars = [tparams[_p(prefix, 'U')],
                   tparams[_p(prefix, 'Ux')]]

    if one_step:
        rval = _step(*(seqs+[init_state, pctx_, pctxx_]+shared_vars))
    else:
        rval, updates = theano.scan(_step,
                                    sequences=seqs,
                                    outputs_info=[init_state],
                                    non_sequences=[pctx_,
                                                   pctxx_]+shared_vars,
                                    name=_p(prefix, '_layers'),
                                    n_steps=nsteps,
                                    profile=profile,
                                    strict=True)
    return rval

# initialize all parameters
def init_params(options):
    params = OrderedDict()
    # embedding
    params['Wemb'] = norm_weight(options['n_words_src'], options['dim_word'])
    params['Wemb_dec'] = norm_weight(options['n_words'], options['dim_word'])
    # encoder
    params = get_layer(options['encoder'])[0](options, params,
                                              prefix='encoder',
                                              nin=options['dim_word'],
                                              dim=options['dim'])
    ctxdim = options['dim']
    # init_state, init_cell
    params = get_layer('ff')[0](options, params, prefix='ff_state',
                                nin=ctxdim, nout=options['dim'])
    # decoder
    params = get_layer(options['decoder'])[0](options, params,
                                              prefix='decoder',
                                              nin=options['dim_word'],
                                              dim=options['dim'],
                                              dimctx=ctxdim)
    # readout
    params = get_layer('ff')[0](options, params, prefix='ff_logit_lstm',
                                nin=options['dim'], nout=options['dim_word'],
                                ortho=False)
    params = get_layer('ff')[0](options, params, prefix='ff_logit_prev',
                                nin=options['dim_word'],
                                nout=options['dim_word'], ortho=False)
    params = get_layer('ff')[0](options, params, prefix='ff_logit_ctx',
                                nin=ctxdim, nout=options['dim_word'],
                                ortho=False)
    params = get_layer('ff')[0](options, params, prefix='ff_logit',
                                nin=options['dim_word'],
                                nout=options['n_words'])

    return params

# build a training model
def build_model(tparams, options):
    opt_ret = dict()

    trng = RandomStreams(1234)
    use_noise = theano.shared(numpy.float32(0.))

    # description string: #words x #samples
    x = tensor.matrix('x', dtype='int64')
    x_mask = tensor.matrix('x_mask', dtype='float32')
    y = tensor.matrix('y', dtype='int64')
    y_mask = tensor.matrix('y_mask', dtype='float32')

    n_timesteps = x.shape[0]
    n_timesteps_trg = y.shape[0]
    n_samples = x.shape[1]

    # word embedding (source)
    emb = tparams['Wemb'][x.flatten()]
    emb = emb.reshape([n_timesteps, n_samples, options['dim_word']])

    # pass through encoder gru, recurrence here
    proj = get_layer(options['encoder'])[1](tparams, emb, options,
                                            prefix='encoder',
                                            mask=x_mask)

    # last hidden state of encoder rnn will be used to initialize decoder rnn
    ctx = proj[0][-1]
    ctx_mean = ctx

    # initial decoder state
    init_state = get_layer('ff')[1](tparams, ctx_mean, options,
                                    prefix='ff_state', activ='tanh')

    # word embedding (target), we will shift the target sequence one time step
    # to the right. This is done because of the bi-gram connections in the
    # readout and decoder rnn. The first target will be all zeros and we will
    # not condition on the last output.
    emb = tparams['Wemb_dec'][y.flatten()]
    emb = emb.reshape([n_timesteps_trg, n_samples, options['dim_word']])
    emb_shifted = tensor.zeros_like(emb)
    emb_shifted = tensor.set_subtensor(emb_shifted[1:], emb[:-1])
    emb = emb_shifted

    # decoder - pass through the decoder gru, recurrence here
    proj = get_layer(options['decoder'])[1](tparams, emb, options,
                                            prefix='decoder',
                                            mask=y_mask, context=ctx,
                                            one_step=False,
                                            init_state=init_state)
    # hidden states of the decoder gru
    proj_h = proj[0]

    # we will condition on the last state of the encoder only
    ctxs = ctx[None, :, :]

    # compute word probabilities
    logit_lstm = get_layer('ff')[1](tparams, proj_h, options,
                                    prefix='ff_logit_lstm', activ='linear')
    logit_prev = get_layer('ff')[1](tparams, emb, options,
                                    prefix='ff_logit_prev', activ='linear')
    logit_ctx = get_layer('ff')[1](tparams, ctxs, options,
                                   prefix='ff_logit_ctx', activ='linear')
    logit = tensor.tanh(logit_lstm+logit_prev+logit_ctx)
    logit = get_layer('ff')[1](tparams, logit, options, prefix='ff_logit',
                               activ='linear')
    logit_shp = logit.shape
    probs = tensor.nnet.softmax(
        logit.reshape([logit_shp[0]*logit_shp[1], logit_shp[2]]))

    # cost
    y_flat = y.flatten()
    y_flat_idx = tensor.arange(y_flat.shape[0]) * options['n_words'] + y_flat
    cost = -tensor.log(probs.flatten()[y_flat_idx])
    cost = cost.reshape([y.shape[0], y.shape[1]])
    cost = (cost * y_mask).sum(0)

    return trng, use_noise, x, x_mask, y, y_mask, opt_ret, cost

# build a sampler
def build_sampler(tparams, options, trng):
    x = tensor.matrix('x', dtype='int64')
    n_timesteps = x.shape[0]
    n_samples = x.shape[1]

    # word embedding (source)
    emb = tparams['Wemb'][x.flatten()]
    emb = emb.reshape([n_timesteps, n_samples, options['dim_word']])

    # encoder
    proj = get_layer(options['encoder'])[1](tparams, emb, options,
                                            prefix='encoder')
    ctx = proj[0][-1]
    ctx_mean = ctx
    init_state = get_layer('ff')[1](tparams, ctx_mean, options,
                                    prefix='ff_state', activ='tanh')

    print 'Building f_init...',
    outs = [init_state, ctx]
    f_init = theano.function([x], outs, name='f_init', profile=profile)
    print 'Done'

    # y: 1 x 1
    y = tensor.vector('y_sampler', dtype='int64')
    init_state = tensor.matrix('init_state', dtype='float32')

    # if it's the first word, emb should be all zero
    emb = tensor.switch(y[:, None] < 0,
                        tensor.alloc(0., 1, tparams['Wemb_dec'].shape[1]),
                        tparams['Wemb_dec'][y])

    # apply one step of gru layer
    proj = get_layer(options['decoder'])[1](tparams, emb, options,
                                            prefix='decoder',
                                            mask=None, context=ctx,
                                            one_step=True,
                                            init_state=init_state)
    next_state = proj
    ctxs = ctx

    # compute the output probability dist and sample
    logit_lstm = get_layer('ff')[1](tparams, next_state, options,
                                    prefix='ff_logit_lstm', activ='linear')
    logit_prev = get_layer('ff')[1](tparams, emb, options,
                                    prefix='ff_logit_prev', activ='linear')
    logit_ctx = get_layer('ff')[1](tparams, ctxs, options,
                                   prefix='ff_logit_ctx', activ='linear')
    logit = tensor.tanh(logit_lstm+logit_prev+logit_ctx)
    logit = get_layer('ff')[1](tparams, logit, options,
                               prefix='ff_logit', activ='linear')
    next_probs = tensor.nnet.softmax(logit)
    next_sample = trng.multinomial(pvals=next_probs).argmax(1)

    # next word probability
    print 'Building f_next..',
    inps = [y, ctx, init_state]
    outs = [next_probs, next_sample, next_state]
    f_next = theano.function(inps, outs, name='f_next', profile=profile)
    print 'Done'

    return f_init, f_next

# generate sample, either with stochastic sampling or beam search
def gen_sample(tparams, f_init, f_next, x, options, trng=None, k=1, maxlen=30,
               stochastic=True, argmax=False):

    # k is the beam size we have
    if k > 1:
        assert not stochastic, \
            'Beam search does not support stochastic sampling'

    sample = []
    sample_score = []
    if stochastic:
        sample_score = 0

    live_k = 1
    dead_k = 0

    hyp_samples = [[]] * live_k
    hyp_scores = numpy.zeros(live_k).astype('float32')
    hyp_states = []

    # get initial state of decoder rnn and encoder context
    ret = f_init(x)
    next_state, ctx0 = ret[0], ret[1]
    next_w = [-1]  # indicator for the first target word (bos target)

    for ii in xrange(maxlen):
        ctx = numpy.tile(ctx0, [live_k, 1])
        inps = [next_w, ctx, next_state]
        ret = f_next(*inps)
        next_p, next_w, next_state = ret[0], ret[1], ret[2]

        if stochastic:
            if argmax:
                nw = next_p[0].argmax()
            else:
                nw = next_w[0]
            sample.append(nw)
            sample_score += next_p[0, nw]
            if nw == 0:
                break
        else:
            cand_scores = hyp_scores[:, None] - numpy.log(next_p)
            cand_flat = cand_scores.flatten()
            ranks_flat = cand_flat.argsort()[:(k-dead_k)]

            voc_size = next_p.shape[1]
            trans_indices = ranks_flat / voc_size
            word_indices = ranks_flat % voc_size
            costs = cand_flat[ranks_flat]

            new_hyp_samples = []
            new_hyp_scores = numpy.zeros(k-dead_k).astype('float32')
            new_hyp_states = []

            for idx, [ti, wi] in enumerate(zip(trans_indices, word_indices)):
                new_hyp_samples.append(hyp_samples[ti]+[wi])
                new_hyp_scores[idx] = copy.copy(costs[idx])
                new_hyp_states.append(copy.copy(next_state[ti]))

            # check the finished samples
            new_live_k = 0
            hyp_samples = []
            hyp_scores = []
            hyp_states = []

            for idx in xrange(len(new_hyp_samples)):
                if new_hyp_samples[idx][-1] == 0:
                    sample.append(new_hyp_samples[idx])
                    sample_score.append(new_hyp_scores[idx])
                    dead_k += 1
                else:
                    new_live_k += 1
                    hyp_samples.append(new_hyp_samples[idx])
                    hyp_scores.append(new_hyp_scores[idx])
                    hyp_states.append(new_hyp_states[idx])
            hyp_scores = numpy.array(hyp_scores)
            live_k = new_live_k

            if new_live_k < 1:
                break
            if dead_k >= k:
                break

            next_w = numpy.array([w[-1] for w in hyp_samples])
            next_state = numpy.array(hyp_states)

    if not stochastic:
        # dump every remaining one
        if live_k > 0:
            for idx in xrange(live_k):
                sample.append(hyp_samples[idx])
                sample_score.append(hyp_scores[idx])

    return sample, sample_score

# calculate the log probablities on a given corpus using translation model
def pred_probs(f_log_probs, prepare_data, options, iterator, verbose=True):
    probs = []

    n_done = 0

    for x, y in iterator:
        n_done += len(x)

        x, x_mask, y, y_mask = prepare_data(x, y,
                                            n_words_src=options['n_words_src'],
                                            n_words=options['n_words'])

        pprobs = f_log_probs(x, x_mask, y, y_mask)
        for pp in pprobs:
            probs.append(pp)

        if numpy.isnan(numpy.mean(probs)):
            ipdb.set_trace()

        if verbose:
            print >>sys.stderr, '%d samples computed' % (n_done)

    return numpy.array(probs)

# optimizers
# name(hyperp, tparams, grads, inputs (list), cost) = f_grad_shared, f_update
def adam(lr, tparams, grads, inp, cost):
    gshared = [theano.shared(p.get_value() * 0.,
                             name='%s_grad' % k)
               for k, p in tparams.iteritems()]
    gsup = [(gs, g) for gs, g in zip(gshared, grads)]

    f_grad_shared = theano.function(inp, cost, updates=gsup, profile=profile)

    lr0 = 0.0002
    b1 = 0.1
    b2 = 0.001
    e = 1e-8

    updates = []

    i = theano.shared(numpy.float32(0.))
    i_t = i + 1.
    fix1 = 1. - b1**(i_t)
    fix2 = 1. - b2**(i_t)
    lr_t = lr0 * (tensor.sqrt(fix2) / fix1)

    for p, g in zip(tparams.values(), gshared):
        m = theano.shared(p.get_value() * 0.)
        v = theano.shared(p.get_value() * 0.)
        m_t = (b1 * g) + ((1. - b1) * m)
        v_t = (b2 * tensor.sqr(g)) + ((1. - b2) * v)
        g_t = m_t / (tensor.sqrt(v_t) + e)
        p_t = p - (lr_t * g_t)
        updates.append((m, m_t))
        updates.append((v, v_t))
        updates.append((p, p_t))
    updates.append((i, i_t))

    f_update = theano.function([lr], [], updates=updates,
                               on_unused_input='ignore', profile=profile)

    return f_grad_shared, f_update

def adadelta(lr, tparams, grads, inp, cost):
    zipped_grads = [theano.shared(p.get_value() * numpy.float32(0.),
                                  name='%s_grad' % k)
                    for k, p in tparams.iteritems()]
    running_up2 = [theano.shared(p.get_value() * numpy.float32(0.),
                                 name='%s_rup2' % k)
                   for k, p in tparams.iteritems()]
    running_grads2 = [theano.shared(p.get_value() * numpy.float32(0.),
                                    name='%s_rgrad2' % k)
                      for k, p in tparams.iteritems()]

    zgup = [(zg, g) for zg, g in zip(zipped_grads, grads)]
    rg2up = [(rg2, 0.95 * rg2 + 0.05 * (g ** 2))
             for rg2, g in zip(running_grads2, grads)]

    f_grad_shared = theano.function(inp, cost, updates=zgup+rg2up,
                                    profile=profile)

    updir = [-tensor.sqrt(ru2 + 1e-6) / tensor.sqrt(rg2 + 1e-6) * zg
             for zg, ru2, rg2 in
             zip(zipped_grads, running_up2, running_grads2)]
    ru2up = [(ru2, 0.95 * ru2 + 0.05 * (ud ** 2))
             for ru2, ud in zip(running_up2, updir)]
    param_up = [(p, p + ud) for p, ud in zip(itemlist(tparams), updir)]

    f_update = theano.function([lr], [], updates=ru2up+param_up,
                               on_unused_input='ignore', profile=profile)

    return f_grad_shared, f_update

def rmsprop(lr, tparams, grads, inp, cost):
    zipped_grads = [theano.shared(p.get_value() * numpy.float32(0.),
                                  name='%s_grad' % k)
                    for k, p in tparams.iteritems()]
    running_grads = [theano.shared(p.get_value() * numpy.float32(0.),
                                   name='%s_rgrad' % k)
                     for k, p in tparams.iteritems()]
    running_grads2 = [theano.shared(p.get_value() * numpy.float32(0.),
                                    name='%s_rgrad2' % k)
                      for k, p in tparams.iteritems()]

    zgup = [(zg, g) for zg, g in zip(zipped_grads, grads)]
    rgup = [(rg, 0.95 * rg + 0.05 * g) for rg, g in zip(running_grads, grads)]
    rg2up = [(rg2, 0.95 * rg2 + 0.05 * (g ** 2))
             for rg2, g in zip(running_grads2, grads)]

    f_grad_shared = theano.function(inp, cost, updates=zgup+rgup+rg2up,
                                    profile=profile)

    updir = [theano.shared(p.get_value() * numpy.float32(0.),
                           name='%s_updir' % k)
             for k, p in tparams.iteritems()]
    updir_new = [(ud, 0.9 * ud - 1e-4 * zg / tensor.sqrt(rg2 - rg ** 2 + 1e-4))
                 for ud, zg, rg, rg2 in zip(updir, zipped_grads, running_grads,
                                            running_grads2)]
    param_up = [(p, p + udn[1])
                for p, udn in zip(itemlist(tparams), updir_new)]
    f_update = theano.function([lr], [], updates=updir_new+param_up,
                               on_unused_input='ignore', profile=profile)

    return f_grad_shared, f_update

def sgd(lr, tparams, grads, x, mask, y, cost):
    gshared = [theano.shared(p.get_value() * 0., name='%s_grad' % k)
               for k, p in tparams.iteritems()]
    gsup = [(gs, g) for gs, g in zip(gshared, grads)]

    f_grad_shared = theano.function([x, mask, y], cost, updates=gsup,
                                    profile=profile)

    pup = [(p, p - lr * g) for p, g in zip(itemlist(tparams), gshared)]
    f_update = theano.function([lr], [], updates=pup, profile=profile)

    return f_grad_shared, f_update

def train(dim_word=100,  # word vector dimensionality
          dim=1000,  # the number of GRU units
          encoder='gru',
          decoder='gru_cond_simple',
          patience=10,  # early stopping patience
          max_epochs=5000,
          finish_after=10000000,  # finish after this many updates
          dispFreq=100,
          decay_c=0.,  # L2 regularization penalty
          alpha_c=0.,  # not used
          lrate=0.01,  # learning rate
          n_words_src=100000,  # source vocabulary size
          n_words=100000,  # target vocabulary size
          maxlen=100,  # maximum length of the description
          optimizer='rmsprop',
          batch_size=16,
          valid_batch_size=16,
          saveto='model.npz',
          validFreq=1000,
          saveFreq=1000,  # save the parameters after every saveFreq updates
          sampleFreq=100,  # generate some samples after every sampleFreq
          datasets=[
              '/data/lisatmp3/chokyun/europarl/europarl-v7.fr-en.en.tok',
              '/data/lisatmp3/chokyun/europarl/europarl-v7.fr-en.fr.tok'],
          valid_datasets=['../data/dev/newstest2011.en.tok',
                          '../data/dev/newstest2011.fr.tok'],
          dictionaries=[
              '/data/lisatmp3/chokyun/europarl/europarl-v7.fr-en.en.tok.pkl',
              '/data/lisatmp3/chokyun/europarl/europarl-v7.fr-en.fr.tok.pkl'],
          use_dropout=False,
          reload_=False):

    # Model options
    model_options = locals().copy()

    # load dictionaries and invert them
    worddicts = [None] * len(dictionaries)
    worddicts_r = [None] * len(dictionaries)
    for ii, dd in enumerate(dictionaries):
        with open(dd, 'rb') as f:
            worddicts[ii] = pkl.load(f)
        worddicts_r[ii] = dict()
        for kk, vv in worddicts[ii].iteritems():
            worddicts_r[ii][vv] = kk

    # reload options
    if reload_ and os.path.exists(saveto):
        with open('%s.pkl' % saveto, 'rb') as f:
            models_options = pkl.load(f)

    print 'Loading data'
    train = TextIterator(datasets[0], datasets[1],
                         dictionaries[0], dictionaries[1],
                         n_words_source=n_words_src, n_words_target=n_words,
                         batch_size=batch_size,
                         maxlen=maxlen)
    valid = TextIterator(valid_datasets[0], valid_datasets[1],
                         dictionaries[0], dictionaries[1],
                         n_words_source=n_words_src, n_words_target=n_words,
                         batch_size=valid_batch_size,
                         maxlen=maxlen)

    print 'Building model'
    params = init_params(model_options)
    # reload parameters
    if reload_ and os.path.exists(saveto):
        params = load_params(saveto, params)

    tparams = init_tparams(params)

    trng, use_noise, \
        x, x_mask, y, y_mask, \
        opt_ret, \
        cost = \
        build_model(tparams, model_options)
    inps = [x, x_mask, y, y_mask]

    print 'Buliding sampler'
    f_init, f_next = build_sampler(tparams, model_options, trng)

    # before any regularizer
    print 'Building f_log_probs...',
    f_log_probs = theano.function(inps, cost, profile=profile)
    print 'Done'

    cost = cost.mean()

    # apply L2 regularization on weights
    if decay_c > 0.:
        decay_c = theano.shared(numpy.float32(decay_c), name='decay_c')
        weight_decay = 0.
        for kk, vv in tparams.iteritems():
            weight_decay += (vv ** 2).sum()
        weight_decay *= decay_c
        cost += weight_decay

    # un used, attention weight regularization
    if alpha_c > 0. and not model_options['decoder'].endswith('simple'):
        alpha_c = theano.shared(numpy.float32(alpha_c), name='alpha_c')
        alpha_reg = alpha_c * (
            (tensor.cast(y_mask.sum(0)//x_mask.sum(0), 'float32')[:, None] -
             opt_ret['dec_alphas'].sum(0))**2).sum(1).mean()
        cost += alpha_reg

    # after all regularizers - compile the computational graph for cost
    print 'Building f_cost...',
    f_cost = theano.function(inps, cost, profile=profile)
    print 'Done'

    print 'Computing gradient...',
    grads = tensor.grad(cost, wrt=itemlist(tparams))
    print 'Done'

    # compile the optimizer, the actual computational graph is compiled here
    lr = tensor.scalar(name='lr')
    print 'Building optimizers...',
    f_grad_shared, f_update = eval(optimizer)(lr, tparams, grads, inps, cost)
    print 'Done'

    print 'Optimization'

    history_errs = []
    # reload history
    if reload_ and os.path.exists(saveto):
        history_errs = list(numpy.load(saveto)['history_errs'])
    best_p = None
    bad_count = 0

    if validFreq == -1:
        validFreq = len(train[0])/batch_size
    if saveFreq == -1:
        saveFreq = len(train[0])/batch_size
    if sampleFreq == -1:
        sampleFreq = len(train[0])/batch_size

    uidx = 0
    estop = False
    for eidx in xrange(max_epochs):
        n_samples = 0

        for x, y in train:
            n_samples += len(x)
            uidx += 1
            use_noise.set_value(1.)

            x, x_mask, y, y_mask = prepare_data(x, y, maxlen=maxlen,
                                                n_words_src=n_words_src,
                                                n_words=n_words)

            if x is None:
                print 'Minibatch with zero sample under length ', maxlen
                uidx -= 1
                continue

            ud_start = time.time()

            # compute cost, grads and copy grads to shared variables
            cost = f_grad_shared(x, x_mask, y, y_mask)

            # do the update on parameters
            f_update(lrate)

            ud = time.time() - ud_start

            # check for bad numbers
            if numpy.isnan(cost) or numpy.isinf(cost):
                print 'NaN detected'
                return 1., 1., 1.

            # verbose
            if numpy.mod(uidx, dispFreq) == 0:
                print 'Epoch ', eidx, 'Update ', uidx, 'Cost ', cost, 'UD ', ud

            # save the best model so far
            if numpy.mod(uidx, saveFreq) == 0:
                print 'Saving...',

                if best_p is not None:
                    params = best_p
                else:
                    params = unzip(tparams)
                numpy.savez(saveto, history_errs=history_errs, **params)
                pkl.dump(model_options, open('%s.pkl' % saveto, 'wb'))
                print 'Done'

            # generate some samples with the model and display them
            if numpy.mod(uidx, sampleFreq) == 0:
                # FIXME: random selection?
                for jj in xrange(numpy.minimum(5, x.shape[1])):
                    stochastic = True
                    sample, score = gen_sample(tparams, f_init, f_next,
                                               x[:, jj][:, None],
                                               model_options, trng=trng, k=1,
                                               maxlen=30,
                                               stochastic=stochastic,
                                               argmax=False)
                    print 'Source ', jj, ': ',
                    for vv in x[:, jj]:
                        if vv == 0:
                            break
                        if vv in worddicts_r[0]:
                            print worddicts_r[0][vv],
                        else:
                            print 'UNK',
                    print
                    print 'Truth ', jj, ' : ',
                    for vv in y[:, jj]:
                        if vv == 0:
                            break
                        if vv in worddicts_r[1]:
                            print worddicts_r[1][vv],
                        else:
                            print 'UNK',
                    print
                    print 'Sample ', jj, ': ',
                    if stochastic:
                        ss = sample
                    else:
                        score = score / numpy.array([len(s) for s in sample])
                        ss = sample[score.argmin()]
                    for vv in ss:
                        if vv == 0:
                            break
                        if vv in worddicts_r[1]:
                            print worddicts_r[1][vv],
                        else:
                            print 'UNK',
                    print

            # validate model on validation set and early stop if necessary
            if numpy.mod(uidx, validFreq) == 0:
                use_noise.set_value(0.)
                valid_errs = pred_probs(f_log_probs, prepare_data,
                                        model_options, valid)
                valid_err = valid_errs.mean()
                history_errs.append(valid_err)

                if uidx == 0 or valid_err <= numpy.array(history_errs).min():
                    best_p = unzip(tparams)
                    bad_counter = 0
                if len(history_errs) > patience and valid_err >= \
                        numpy.array(history_errs)[:-patience].min():
                    bad_counter += 1
                    if bad_counter > patience:
                        print 'Early Stop!'
                        estop = True
                        break

                if numpy.isnan(valid_err):
                    ipdb.set_trace()

                print 'Valid ', valid_err

            # finish after this many updates
            if uidx >= finish_after:
                print 'Finishing after %d iterations!' % uidx
                estop = True
                break

        print 'Seen %d samples' % n_samples

        if estop:
            break

    if best_p is not None:
        zipp(best_p, tparams)

    use_noise.set_value(0.)
    valid_err = pred_probs(f_log_probs, prepare_data,
                           model_options, valid).mean()

    print 'Valid ', valid_err

    params = copy.copy(best_p)
    numpy.savez(saveto, zipped_params=best_p,
                history_errs=history_errs,
                **params)

    return valid_err

if __name__ == '__main__':
    pass