#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import pickle	#For read inputs fields
import bson
import sys		#For sys.exit
import numpy	#For handle inputs fields and output fields
import schromod	#Our C++ binding for solving stuff

from pymongo import MongoClient
import pymongo.errors

from evtk.hl import imageToVTK

#import matplotlib.pyplot as plt

pickle_path = "../../generator/"

username, password, host, dbname = 'tazio', 'taziotazio', '127.0.0.1', 'CarTata'
client = MongoClient('mongodb://%s:%s@%s/%s' % (username, password, host, dbname))

current_run_id = int(input("Run id ? ")) 


data_input = client[dbname]['data_input'].find_one({'run_id':current_run_id})
potential_u = pickle.loads(data_input['field_potential'])
potential = numpy.array(potential_u,dtype=numpy.float64,order = 'C')

mesh = data_input['mesh']
nx = mesh['width']
ny = mesh['length']


print("Saving potential in potential.vti")
imageToVTK("data/potential", pointData = {'V': numpy.reshape(potential,(nx,ny,1)).copy()})

#We don't take iteration == 0 cause it's not normalized and it disturb the plot
run_data_i = client[dbname]['data_psi'].find({'run_id':current_run_id, "iteration": {"$gt":0}}, sort = [("iteration", 1)])

#progress bar initialization
bar_width = 60
bar_offset = 30
sys.stdout.write("%s[%s]" % (" " * bar_offset,  "⋅" * bar_width))
sys.stdout.flush()
sys.stdout.write("\b" * (bar_width+1)) # return to start of line, after '['
true_width = run_data_i.count() / bar_width
progress_offset = 1

for i, data in enumerate(run_data_i):
    psi_u = pickle.loads(data['mat'])
    psi = numpy.array(psi_u,dtype=numpy.complex128,order = 'C')

    filename = 'data/anim_%07d' % (data['iteration'])
    file_str = filename + " created"

    neww = numpy.absolute(psi)
    neww = neww.reshape((nx,ny,1), order='C').copy()

    real = numpy.real(psi)
    real = real.reshape((nx,ny,1), order='C').copy()

    imag = numpy.imag(psi)
    imag = real.reshape((nx,ny,1), order='C').copy()
    imageToVTK(filename, pointData = {'psi_norm': neww, 'real': real, 'imag': imag})
    sys.stdout.write("\r"*(len(file_str)+1) + file_str)
	
    #progress bar +1
    if not i % true_width:
        sys.stdout.write('\r\033[%sC#' % (str(bar_offset + progress_offset)))
	progress_offset %= bar_width
	progress_offset += 1
    sys.stdout.flush()

sys.stdout.write("\n")