# -*- coding: utf-8 -*-

import struct
import numpy as np
import multiprocessing # To get CPU count.
import os
import tempfile
import subprocess

import phstuff.barcode as bc
import phstuff.simplicial as simpl

"""Handle marshalling data into and out of DIPHA's formats, and
optionally manage running the DIPHA executable.

"""

# FIXME, TODO: These magic numbers should really be read from the
# DIPHA header files.
DIPHA_MAGIC = 8067171840
DIPHA_WEIGHTED_BOUNDARY_MATRIX = 0
DIPHA_IMAGE_DATA = 1
DIPHA_PERSISTENCE_DIAGRAM = 2
DIPHA_DISTANCE_MATRIX = 7
DIPHA_SPARSE_DISTANCE_MATRIX = 8


def save_weight_matrix(fname, weights):
    """Write a NumPy weight matrix to a DIPHA full distance matrix
    file. Corresponds to DIPHA's "distance matrix data".

    Attention: DIPHA currently uses a convention wherein an edge's
    filtration value is *half* its weight. Prior to version 0.0.6 of
    this package, we would compensate for this by multiplying every
    weight by 2. As of version 0.0.6, we no longer do so. Pay
    attention in case DIPHA's behavior changes in the future!

    Parameters:
    -----------

    fname: Name of file to write.

    weights: NumPy array. Elements will be converted to IEEE-754
       doubles and used as weights. The entire array is passed on to
       DIPHA, whose behvior is undefined if it's not symmetric.

    """

    # override_dipha_half: DIPHA currently (commit 0081862) divides all
    #    entries by 2 when loading the file. If this argument is `True`,
    #    we compensate by multiplying all entries by 2. Be sure to pay
    #    attention if DIPHA's behavior changes!

    factor = 1.0
    
    m = weights.shape[0]
    if weights.shape[0] != weights.shape[1]:
        raise ValueError("Matrix is not square.")

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_DISTANCE_MATRIX))
        f.write(struct.pack("<q", m))

        # Why write in a loop? I heard rumors that struct.pack will
        # allocate a lot of temporary heap space if given lots of
        # data. I haven't actually tested.
        for i in range(0, m):
            for j in range(0, m):
                f.write(struct.pack("<d", factor*weights[i,j]))


def load_weight_matrix(fname):
    """Load a DIPHA full distance matrix as a NumPy array.

    Parameters:
    -----------

    fname: Name of file to read.

    Returns:
    ----------

    A NumPy array.

    """

    with open(fname, "rb") as f:
        if struct.unpack('<q', f.read(8))[0] != DIPHA_MAGIC:
            raise IOError("File %s is not a valid DIPHA file." %(fname))
        if struct.unpack('<q', f.read(8))[0] != DIPHA_DISTANCE_MATRIX:
            raise IOError("File %s is not a valid DIPHA matrix file." %(fname))

        m = struct.unpack('<q', f.read(8))[0]
        if m <= 0:
            raise IOError("Matrix dimension non-sensical (%d)." %(m))

        X = np.zeros((m, m))

        for i in range(0, m):
            for j in range(0, m):
                X[i, j] = struct.unpack('<d', f.read(8))[0]

    return X
    

def save_masked_weight_matrix(fname, weights):
    """Write a masked NumPy weight matrix to a DIPHA sparse distance
    matrix file, keeping edges only for unmasked entries. Corresponds
    to DIPHA's "sparse distance matrix data".

    Attention: DIPHA currently uses a convention wherein an edge's
    filtration value is *half* its weight. Prior to version 0.0.6 of
    this package, we would compensate for this by multiplying every
    weight by 2. As of version 0.0.6, we no longer do so. Pay
    attention in case DIPHA's behavior changes in the future!

    Parameters:
    -----------

    fname: Name of file to write.

    weights: Masked NumPy array. Unmasked elements will be converted
       to IEEE-754 and used as weights. Must be symmetric (not
       checked).

    """

    factor = 1.0

    n = weights.shape[1]

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_SPARSE_DISTANCE_MATRIX))
        f.write(struct.pack("<q", n))

        for i in range(0, n):
            f.write(struct.pack('<q', n - np.sum(weights.mask[i, :])))
        for i in range(0, n):
            for j in np.arange(0, n)[-(weights.mask[i, :])]:
                    f.write(struct.pack('<q', j))
        for i in range(0, n):
            for j in np.arange(0, n)[-(weights.mask[i, :])]:
                    f.write(struct.pack('<d', factor*weights[i, j]))
    


def save_edge_list(fname, edge_list):
    """Write a possibly non-complete weighted graph represented as an edge
    list to a DIPHA sparse distance file. Corresponds
    to DIPHA's "sparse distance matrix data".

    Attention: DIPHA currently uses a convention wherein an edge's
    filtration value is *half* its weight. Prior to version 0.0.6 of
    this package, we would compensate for this by multiplying every
    weight by 2. As of version 0.0.6, we no longer do so. Pay
    attention in case DIPHA's behavior changes in the future!

    Parameters:
    -----------

    fname: Name of file to write.

    edge_list: A list of the edges and weights of each node, in node
       order. If node `i` has edges to nodes `a_{i,0}, …¸ a_{i,k_i}`
       with weights `w_{i,0}, …, w_{i,k_i}`, then this argument should
       be

          `[[(a_{0,0}, w_{0,0}), …, (a_{0,k_0}, w_{0,k_0})],
            [(a_{1,0}, w_{1,0}), …, (a_{1,k_1}, w_{1,k_1})], 
            …, 
            [(a_{n,0}, w_{n,0}), …, (a_{n,k_n}, w_{n,k_n})]]`

       Remember that an isolated node corresponds to an empty (inner)
       list.

    """

    factor = 1.0
    
    n = len(edge_list)

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_SPARSE_DISTANCE_MATRIX))
        f.write(struct.pack("<q", n))

        for i in range(0, n):
            f.write(struct.pack('<q', len(edge_list[i])))
        for i in range(0, n):
            for j in range(0, len(edge_list[i])):
                f.write(struct.pack('<q', edge_list[i][j][0]))
        for i in range(0, n):
            for j in range(0, len(edge_list[i])):
                f.write(struct.pack('<d', factor*edge_list[i][j][1]))

def save_cubical(fname, array):
    """Save a (any-dimensional) array that will be interpreted as a
    cubical complex with top-dimensional cells entering the filtration
    at the scale given by the array values.

    Corresponds to DIPHA's "image data".

    Parameters:
    -----------
    
    fname: Name of file to write.

    array: Any-dimensional array with floating-point values giving
    the filtration values of the top-dimensional cells.

    """

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_IMAGE_DATA))
        f.write(struct.pack("<q", np.product(array.shape)))
        f.write(struct.pack("<q", len(array.shape)))
        
        for n in array.shape:
            f.write(struct.pack("<q", n))

        flat = array.flattened(order="F")
        for w in flat:
            f.write(struct.pack("<d", float(w)))
            
def save_simplicial(fname, complex):
    """*EXPERIMENTAL*. Save a filtered simplicial complex.

    Corresponds to DIPHA's "weighted boundary matrix".

    Parameters:
    -----------
    
    fname: Name of file to write.

    complex: Simplicial complex to save.

    """
        
    complex.order()
    assert(complex.is_ordered())
    
    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_WEIGHTED_BOUNDARY_MATRIX))
        f.write(struct.pack("<q", 0))
        f.write(struct.pack("<q", complex.size()))
        f.write(struct.pack("<q", complex.top_dim()))

        # This can be made a whole lot more efficient. We're iterating
        # in a stupid way.
        
        offsets = []
        total_nonzero = 0
        for node in complex:
            d = node.dimension()
            f.write(struct.pack("<q", d))
            offsets.append(total_nonzero)
            if d != 0:
                total_nonzero += d+1                

        for node in complex:
            f.write(struct.pack("<d", node.w))            

        for off in offsets:
            f.write(struct.pack("<q", off))

        f.write(struct.pack("<q", total_nonzero))
        for node in complex:
            for bnode in node.boundary_nodes():
                f.write(struct.pack("<q", bnode.i))
            
            
def load_barcode(fname, top_dim = None):

    """Load a barcode as written by DIPHA.

    Parameters:
    -----------

    fname: File name to load.
    
    top_dim: Integer. Exclude degrees above this.

    Returns:
    --------

    A dict keyed on persistent homology degree. Associated to each
    degree is a list of `Interval`s.

    """
    
    ret = dict()
    
    with open(fname, "rb") as f:
        if struct.unpack('<q', f.read(8))[0] != DIPHA_MAGIC:
            raise IOError("File %s is not a valid DIPHA file." %(fname))
        if struct.unpack('<q', f.read(8))[0] != DIPHA_PERSISTENCE_DIAGRAM:
            raise IOError("File %s is not a valid DIPHA barcode file." %(fname))

        n = struct.unpack('<q', f.read(8))[0]
        
        for i in range(0, n):
            (d, birth, death) = struct.unpack('<qdd', f.read(3*8))
            if d < 0:
                dim = -d -1
            else:
                dim = d

            if top_dim is None or dim <= top_dim:
                if d < 0:
                    ret.setdefault(dim, []).append(bc.Interval(birth))
                else:
                    ret.setdefault(dim, []).append(bc.Interval(birth, death))

    return ret

def save_text_barcode(fname, barcode):
    """Saves a barcode as a human- and machine-readable text file in an
    ad-hoc format.

    Parameters:
    -----------
    
    fname: File name to save to.
    
    barcode: The barcode to save, a dictionary keyed on degree, each
    entry being a list of `Interval`s.

    """

    with open(fname, "w") as f:
        for dim in barcode.keys():
            for interval in barcode[dim]:
                if interval.is_finite():
                    f.write("%d %g %g\n" %(dim, interval.birth, interval.death))
                else:
                    f.write("%d %g inf\n" %(dim, interval.birth))

def save_barcode(fname, barcode):
    """Saves a barcode in DIPHA's format.

    Parameters:
    -----------
    
    fname: File name to save to.
    
    barcode: The barcode to save, a dictionary keyed on degree, each
    entry being a list of `Interval`s.

    """

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_PERSISTENCE_DIAGRAM))

        n = 0
        for (dim, intervals) in barcode.items():
            n += len(intervals)
        f.write(struct.pack("<q", n))
        
        for (dim, intervals) in barcode.items():
            for interval in intervals:
                if interval.is_finite():
                    f.write(struct.pack("<qdd", dim, interval.birth, interval.death))
                else:
                    f.write(struct.pack("<qdd", -dim - 1, interval.birth, float("inf")))

    

class DiphaRunner:
    """For long-running computations that may fail, or computations
    involving huge amounts of data, I recommend that you run DIPHA
    yourself. If you prefer to avoid doing that, this class helps you
    run DIPHA without leaving your Python script.

    In order to do so, the `DiphaRunner` object needs to know where to
    find the `mpirun` and `dipha` executables. If they are somewhere
    in your environment's `PATH`, then everything is fine. If not, you
    need to specify where these can be foind either by setting the
    environment variables `DIPHA`/`MPIRUN` (to the full path of the
    executables, not just their directories) or through the optional
    arguments of the class constructor.

    """

    def __init__(self, top_dim, cpu_count = None, quiet = True, mpirun = None, dipha = None):
        """Constructor.
        
        Parameters:
        -----------
        
        top_dim: Highest simplex dimension to include.

        cpu_count (optional): The number of parallell DIPHA processes
        to run. If `None` (default), try to autodetect. The actual
        number that will be used can be found in the `cpu_count`
        member of the returned object.

        quiet (optional): Don't print anything in normal situations
        (default).

        mpirun (optional): Full path of the `mpirun` executable. If
        `None` (default), use environment variable `MPIRUN` or look in
        `PATH`.

        dipha (optional): Full path of the `dipha` executable. If
        `None` (default), use environment variable `DIPHA` or look in
        `PATH`.

        """
    
        self.top_dim = int(top_dim)
        if top_dim <= 0:
            raise ValueError("Top dimension must be strictly positive.")
        
        if cpu_count is None:
            try:
                self.cpu_count = multiprocessing.cpu_count()
            except NotImplementedError as e:
                self.cpu_count = 1
        else:
            self.cpu_count = cpu_count

        self.quiet = quiet

        if mpirun is None:
            self.mpirun = os.environ.get("MPIRUN")
            paths = os.environ.get("PATH", "")
            if self.mpirun is None:
                for path in paths.split(":"):
                    x = os.path.join(path, "mpirun")
                    if os.path.isfile(x):
                        self.mpirun = x
                        break
        else:
            self.mpirun = mpirun

        if dipha is None:
            self.dipha = os.environ.get("DIPHA")
            paths = os.environ.get("PATH", "")
            if self.dipha is None:
                for path in paths.split(":"):
                    x = os.path.join(path, "dipha")
                    if os.path.isfile(x):
                        self.dipha = x
                        break
        else:
            self.dipha = dipha

        if self.mpirun is None or not os.path.isfile(self.mpirun):
            raise RuntimeError("Could not find mpirun exectuable.")
        if self.dipha is None or not os.path.isfile(self.dipha):
            raise RuntimeError("Could not find dipha executable.")

        self.tmpdir = tempfile.mkdtemp()
        self.ifile = os.path.join(self.tmpdir, "input.dipha")
        self.ofile = os.path.join(self.tmpdir, "output.dipha")
        f = open(self.ifile, "w")
        f.close()
        f = open(self.ofile, "w")
        f.close()
        self.ready = False

        
    def weight_matrix(self, weights):
        """Initialize to run with a weight matrix. See `save_weight_matrix`."""
        
        save_weight_matrix(self.ifile, weights)
        self.ready = True

    def masked_weight_matrix(self, weights):
        """Initialize to run with a masked weight matrix. See
        `save_masked_weight_matrix`."""
        
        save_masked_weight_matrix(self.ifile, weights)
        self.ready = True

    def edge_list(self, edge_list):
        """Initialize to run with an edge list. See `save_edge_list`."""
        
        save_edge_list(self.ifile, edge_list)
        self.ready = True

    def cubical(self, array):
        """Initialize to run with a cubical complex. See `save_cubical`."""

        save_cubical(self.ifile, array)
        self.ready = True

    def simplicial(self, complex):
        save_simplicial(self.ifile, complex)
        self.ready = True

    def run(self):
        """Run DIPHA. Success is indicated by `code` member of the object
        being zero. If that is the case, then the result is available
        in the `barcode` member."""

        if not self.ready:
            self.cleanup()
            raise RuntimeError("Not ready to run. Maybe you haven't given me data, or I've already been run?")
        
        if self.quiet:
            opipe = subprocess.PIPE
        else:
            opipe = None

        if not self.quiet:
            print("Spawning %d DIPHA processes." %(self.cpu_count))
            
        proc = subprocess.Popen([self.mpirun, "-np", "%d" %(self.cpu_count), self.dipha, "--upper_dim", "%d" %(self.top_dim), self.ifile, self.ofile], stdout=opipe)

        self.out_message = proc.communicate()[0]
        self.code = proc.returncode
        
        if not self.quiet and self.code != 0:
            print("DIPHA exited abnormally with code %d." %(self.code))
            self.cleanup()

        if self.code == 0:
            self.barcode = load_barcode(self.ofile, top_dim = self.top_dim - 1)
            if not self.quiet:
                print("DIPHA exited normally. Result available.")

        self.ready = False
        self.cleanup()
        

    def cleanup(self):
        os.remove(self.ifile)
        os.remove(self.ofile)
        os.rmdir(self.tmpdir)
        self.ifile = None
        self.ofile = None
        self.tmpdir = None