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

import hashlib
import bson
import pickle	#For read inputs fields
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


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


######################
#If you want ot change methods : it's here
def run(a):
	#a.FTCS()
	#a.BTCS(50)
	a.CTCS(50)
######################

def norm_mesh(arr,hx,hy):
	return numpy.absolute(numpy.sum(numpy.multiply(numpy.conj(arr),arr))/(hx*hy))

def normalize(arr,hx,hy):
	return numpy.divide(arr,numpy.float64(norm_mesh(arr,hx,hy)))

def hash_data(content):
	m = hashlib.md5()
	m.update(content)
	return m.hexdigest()

#Used after a.V is initialized
def check_n_restart(a,iteration,psi):
	if iteration != 0:
		a.setPsiInit(psi)
	else:
		a.setPsiInit(normalize(psi,a.hx,a.hy))

current_run_id = int(input("Run id ? ")) 
data_input = client[dbname]['data_input'].find_one({'run_id':current_run_id})
try:
    bin_data=data_input['field_potential']
except TypeError:
    print("You have to choose a valid run_id")
    exit()

if hash_data(bin_data) != data_input['field_potential_bin_hash']:
	print('ERROR : no matching hash for potential')
	exit()

potential_u = pickle.loads(bin_data)
potential = numpy.array(potential_u,dtype=numpy.float64,order = 'F')
run_data = client[dbname]['data_psi'].find_one({'run_id':current_run_id}, sort = [("iteration", -1)])
try:
    bin_data=run_data['mat']
except TypeError:
    print("You have to choose a valid run_id")
    exit()

if hash_data(bin_data) != run_data['mat_bin_hash']:
	print('ERROR : not matching hash for psi')
	exit()
psi_u = pickle.loads(bin_data)
psi = numpy.array(psi_u,dtype=numpy.complex128,order = 'F')

nb_iter = data_input['nb_iter']
dtime = data_input['dtime']
mesh = data_input['mesh']
img_factor = data_input['img_fact']
iteration = run_data["iteration"] 

print("Start from last checkpoint %d" % (iteration*img_factor))

a = schromod.Solveur()
a.setV(potential)

a.ht=dtime
a.hx=mesh['stitch']['width']
a.hy=mesh['stitch']['length']

check_n_restart(a,iteration,psi)
print('Psi is well initalized')
print("psi norm : %f"%(norm_mesh(a.psi,a.hx,a.hy)))
print('Begin of calculation')

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

if iteration != 0:
	run(a)
	psi_bin = bson.binary.Binary(pickle.dumps(a.getPsi(), protocol=2))
	client[dbname]['data_psi'].replace_one({'run_id':current_run_id, 'iteration':iteration}, 
						{'run_id':current_run_id, 'mat':psi_bin, 'iteration':iteration})

for i in range((iteration * img_factor), nb_iter):
    run(a)

    if i % img_factor == 0:
        print("psi norm at %d step : %f"%(i,norm_mesh(a.getPsi(),a.hx,a.hy)))
        iteration += 1
        psi_bin = bson.binary.Binary(pickle.dumps(a.getPsi(), protocol=2))
        hash_m=hash_data(psi_bin)
        client[dbname]['data_psi'].insert_one({'mat':psi_bin, 'mat_bin_hash':hash_m,
                                              'run_id':current_run_id, 'iteration':iteration})