# -*- coding: utf-8 -*-
# Copyright (C) 2008 Kent-Andre Mardal
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function
import sys
import getopt
import six
from dolfin import *
import numpy
import operator
from functools import reduce

def create_maps(boundary_mesh):
    """
    Create mappings between vertices and cells on the boundary mesh
    """

    c2v = {}
    v2c = {}
    for cell in cells(boundary_mesh):
        for vertex in vertices(cell):
            v = vertex.index()
            c = cell.index()
            if c in c2v:
                c2v[c].append(v)
            else:
                c2v[c] = [v]
            if v in v2c:
                v2c[v].append(c)
            else:
                v2c[v] = [c]
    return v2c, c2v


def find_subdomain(start_vertex, v2c, c2v):
    """
    Find all cells connected with the start_vertex
    """
    next  = c2v[v2c[start_vertex][0]][1]

    done = False
    previous_vertex = start_vertex
    previous_cell   = v2c[start_vertex][0]
    subdomain = {}
    while not done:
        c0, c1 = v2c[previous_vertex]
        if c0 == previous_cell:
            next_cell = c1
        elif c1 == previous_cell:
            next_cell = c0
        else:
            print("None of the vertices were the previous vertex ??")
        v0, v1 = c2v[next_cell]
        if v0 == previous_vertex:
            next_vertex = v1
        elif v1 == previous_vertex:
            next_vertex = v0
        else:
            print("None of the vertices were the previous vertex ??")
        if next_cell in subdomain:
            done = True
        else:
            subdomain[next_cell] = True
            previous_vertex = next_vertex
            previous_cell   = next_cell
    return subdomain

def find_new_cell_key(c2v, keys):
    """
    Find new cell.
    """
    done = False
    iter = list(c2v.keys()).__iter__()
    while not done:
        key = six.advance_iterator(iter)
        if key in keys:
            pass
        else:
            done = True
    return key


def write_file(outfilename, mesh, subdomains):
    str = """<?xml version="1.0" encoding="UTF-8"?>
<dolfin xmlns:dolfin="http://fenicsproject.org">
  <meshfunction type="uint" dim="1" size="%d">
    """ % mesh.numFacets()

    for key in list(subdomains.keys()):
        str += "\n     <entity index=\"%d\" value=\"%d\"/>" % (key, subdomains[key])
    str +=  """
  </meshfunction>
</dolfin>
    """

    f = open(outfilename, 'w')
    f.write(str)
    f.close()


def find_keys_on_one_subdomain(c2v, v2c, all_keys):
    try:
        flat_keys = []
        if len(all_keys) > 0:
            flat_keys = reduce(operator.add, all_keys)

        new_cell_key = find_new_cell_key(c2v, flat_keys)
        vertex_key = c2v[new_cell_key][0]
        subdomain = find_subdomain(vertex_key, v2c, c2v)
        new_keys = list(subdomain.keys());
        new_keys.sort()
        all_keys.append(new_keys)
        return all_keys
    except Exception as e:
#        print "done with finding the subdomains"
        pass




def find_all_subdomains(ifilename, ofilename):

    mesh = Mesh(ifilename)

    if mesh.geometry().dim() != 2:
        print("This script does only work in 2D.")
        print("(It should be possible to extend to 3D, though)")
        exit(2)

    boundary_mesh = BoundaryMesh(mesh)

    boundary_cell_map = boundary_mesh.data().meshFunction("cell map")
    v2c, c2v = create_maps(boundary_mesh)

    done = False
    keys = []
    subdomains = {}
    all_keys = []
    prev_keys = []
    while not done:
        new_keys = find_keys_on_one_subdomain(c2v, v2c, prev_keys)
        if new_keys == None:
            done = True
            all_keys = prev_keys
        else:
            prev_keys = new_keys

    for i in range(0, mesh.numFacets()):
        subdomains[i] = len(all_keys)

    counter = 0
    for keys in all_keys:
        for key in keys:
            subdomains[boundary_cell_map.array()[key]] =  counter
        counter+=1

    write_file(ofilename, mesh, subdomains)


def usage():
    "Display usage"
    print("""\
Usage: dolfin-convert [OPTIONS] ... input.x output.y

Options:

  -h         display this help text and exit
  -i specify input file
  -o specify output file

Alternatively, the following long options may be used:

  --help     same as -h
  --input    same as -i
  --output   same as -o

""")



def main(argv):
    "Main function"
         # Get command-line arguments
    try:
        opts, args = getopt.getopt(argv, "hi:o:", ["help", "input=", "output="])
    except getopt.GetoptError:
        usage()
        sys.exit(2)


    ifilename = None
    ofilename = None
    for opt, arg in opts:
        if opt in ("-i", "--infile"):
            ifilename = arg
        elif opt in ("-o", "--outfile"):
            ofilename = arg
        else:
            print("Option not recoginized")

    find_all_subdomains(ifilename, ofilename)





if __name__ == '__main__':
    main(sys.argv[1:])




