#!/usr/bin/env python

import scipy.sparse

import numpy as np

import graph_tool.all as gt

from numpy.random import binomial

import logging as log

import petsc4py

import sys

#petsc4py.init(sys.argv)

from petsc4py import PETSc

import argparse

import resnet

parser = argparse.ArgumentParser(description='Cubic resistor lattice solver using PETSc. Specify either an npy file for the graph edge filter or generate one randomly for the specified bond probability.')

parser.add_argument('N', type=int, default=100, help='number of lattice sites on an edge, needs to be specified if loading a filter or generating one')

parser.add_argument('--gm', type=float, default=1.0, help='bond conductance')

group = parser.add_mutually_exclusive_group(required=True)

group.add_argument('--maskfile',type=str)

group.add_argument('-p','--probability',type=float,metavar='P',help='probability of placing a bond')

args = parser.parse_args()

log.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=log.INFO)

class MyVisitor(gt.DFSVisitor):

self.visited.add(int(u))

Lx = args.N

Ly = Lz = Lx

l = Lx*Ly*Lz

rm = 1.0

gm = args.gm

num_bonds = (Lx-1)*Ly*Lz + Lx*(Ly-1)*Lz + Lx*Ly*(Lz-1)

log.info("generating graph")

g = gt.lattice([Lx,Ly,Lz])

assert(g.num_edges() == num_bonds)

p = args.probability

if args.maskfile:

log.info("loading %s" % args.maskfile)

mask = np.load(args.maskfile)

p = np.sum(mask)/mask.shape[0]

vals = mask

else:

vals = binomial(1,p,num_bonds).astype('bool')

prop = g.new_edge_property('bool',vals=vals)

g.set_edge_filter(prop)

num_vert = g.num_vertices()

g.add_vertex(2)

g.add_edge_list([[num_vert,n] for n in range(Lx*Ly)])

g.add_edge_list([[num_vert+1,n] for n in range(l-Lx*Ly,l)])

log.info("checking for percolation")

vcc_path = set()

gt.dfs_search(g,num_vert,visitor=MyVisitor(vcc_path))

if not num_vert+1 in vcc_path:

log.warn("not percolating")

else:

vcc_path.remove(num_vert+1)

vcc_path.remove(num_vert)

g.remove_vertex(num_vert+1)

g.remove_vertex(num_vert)

plane = [i for i in range(Lx*Ly,2*Lx*Ly+1) if i in g.vertex(i-Lx*Ly).all_neighbours() and i in vcc_path]

top = range(Lx*Ly)

bot = range(l-Lx*Ly,l)

others = vcc_path-set(top)-set(bot)

fixed = set(range(l))-others

adj = gt.adjacency(g)

del g

resnet.csr_rows_set_nz_to_val(adj,fixed) # nodes not in vcc_path or those on the top/bottom

diag = adj.sum(axis=1)

log.info("creating PETSc structures")

b = PETSc.Vec().createSeq(l)

x = PETSc.Vec().createSeq(l)

#A = PETSc.Mat().createAIJ(size=adj.shape,nnz=7,csr=(adj.indptr,adj.indices,-adj.data))

A = PETSc.Mat().createAIJ(size=adj.shape,nnz=7)

# the gm factor here and below doesn't actually matter

# scaling each resistor by the same constant still produces

# the same distribution of voltages

A.setValuesCSR(adj.indptr,adj.indices,-gm*adj.data)

for i in fixed:

A.setValue(i,i,1.0)

for i in others:

A.setValue(i,i,gm*diag[i])

del adj

b.setValuesBlocked(range(Lx*Ly),[1.0]*Lx*Ly)

# i don't know if this actually frees memory

#del g

log.info("assembling")

A.assemblyBegin()

A.assemblyEnd()

ksp = PETSc.KSP().create()

ksp.setOperators(A)

# uncomment if you want direct LU instead of iterative

# ksp.setType('preonly')

# pc=ksp.getPC()

# pc.setType('lu')

# pc.setFactorSolverPackage('mumps')

ksp.setFromOptions()

log.info("solving with %s" % ksp.getType())

ksp.solve(b,x)

log.info("converged in %d iterations" % ksp.getIterationNumber())

V = x.array

Itot = np.sum([(1.0-V[i])*gm for i in plane])

log.info("total current %e" % Itot)

#emt = gm*(1.-(1.-p)*3./2.)/(args.N-1)*(args.N)**2

emt = 0.125*(-2*gm+6*gm*p+np.sqrt(gm**2*(-2+6*p)**2))/(args.N-1)*args.N**2

log.info("emt prediction %e" % emt)