#!/usr/bin/env python3

# Lappy is intended to be a drop-in replacement for Lapdog in cases
# where one prefers simple dependencies over high performance and
# scalability.
#
# One incompatibility is obviously that Lappy does not respect
# PETSc/SLEPc command line switches, and one therefore has to specify
# the number of eigenvalues/eigenvectors through --number.

import scipy.sparse.linalg
import sys
import numpy as np
import petsc

def print_usage():
    sys.stderr.write("Usage:\n")
    sys.stderr.write("%s --laplacian|-l file --number|-n n --vals file --vecs file [-h|--help]\n" %(sys.argv[0]))

def main():
    n = 0
    i = 1
    while i < len(sys.argv):
        arg = sys.argv[i]
        if arg in ["-h", "--help"]:
            print_usage()
            exit(0)
        elif arg in ["-l", "--laplacian"]:
            if i + 1 < len(sys.argv) and sys.argv[i+1][0] != "-":
                infile = sys.argv[i+1]
                i += 1
            else:
                sys.stderr.write("Missing argument to --laplacian.\n")
                print_usage()
                exit(1)
        elif arg in ["--vals"]:
            if i + 1 < len(sys.argv) and sys.argv[i+1][0] != "-":
                outfile_vals = sys.argv[i+1]
                i += 1
        elif arg in ["--vecs"]:
            if i + 1 < len(sys.argv) and sys.argv[i+1][0] != "-":
                outfile_vecs = sys.argv[i+1]
                i += 1
        elif arg in ["-n", "--number"]:
            if i + 1 < len(sys.argv):
                try:
                    n = int(sys.argv[i+1])
                    i += 1
                except ValueError as err:
                    sys.stderr.write("Invalid argument to --number.\n")
                    print_usage()
                    exit(1)
        else:
            sys.stderr.write("Unknown option %s.\n" %(arg))
            print_usage()
            exit(1)
        i += 1

    if infile == "" or outfile_vals == "" or outfile_vecs == "":
        sys.stderr.write("Need a Laplacian and both input files.\n")
        print_usage()
        exit(1)
    if n <= 0:
        sys.stderr.write("The --number parameter is mandatory and must be >0.\n")
        print_usage()
        exit(1)
            
    L = petsc.load(infile)
    print("Laplacian is size %d." %(L.shape[0]))
    
    (vals, vecs) = scipy.sparse.linalg.eigsh(L, k=min(n, L.shape[0]), M=None, sigma=None, which="SM", v0=None, ncv=None, maxiter=None, tol=0, return_eigenvectors=True, Minv=None, OPinv=None, mode="normal")

    print("Have %d eigenpairs." %(len(vals)))
    cutoff = 1e-6 # Arbitrary
    estimated_betti = 0
    with open(outfile_vals, "w") as f:
        for i in range(0, len(vals)):
            f.write("%f\n" %(vals[i]))
            print("%d: %.20f" %(i, vals[i]))
            if np.abs(vals[i]) < cutoff:
                estimated_betti += 1

    print("Estimated Betti number: %d." %(estimated_betti))

    np.savetxt(outfile_vecs, vecs.T)

if  __name__ == "__main__":
    main()