import numpy as np
import pyray.shapes.twod.square_mesh as sm
import pyray.shapes.fourd.tesseract_graph as tg
from copy import deepcopy
import json
from PIL import Image, ImageDraw
import pyray.shapes.fourd.tst_open_tsrct as ot
from pyray.rotation import rotation
import os
from collections import defaultdict


def refresh_mesh(tf):
    tf.reset_rotations()
    tf.dfs_flatten2('00-+')
    tf.reset_vert_col()
    tf.actually_flatten()


def gen_tsrct_meshes(num_mshs=100):
    tf = tg.TsrctFcGraph(angle=np.pi/2)
    tf.dfs_flatten2('00-+')
    tf.reset_vert_col()
    tf.actually_flatten()
    msh_files = os.listdir("Data/Meshes")
    msh_files = sorted(msh_files)
    print("Number of mesh files before: " + str(len(msh_files)))
    base_i = 0
    try:
        base_i = int(msh_files[len(msh_files)-1].replace('mesh','').replace('.txt',''))
    except:
        pass
    ii = 0
    while ii < num_mshs:
        # First get the two meshes.
        msh = sm.SqMesh(tf)
        f1 = np.random.choice([k for k in tf.adj.keys()])
        f2 = np.random.choice([k for k in tf.adj[f1]])
        #f1 = '-0-0'
        #f2 = '00-+'
        m1, m2 = msh.cut_mesh(f1, f2)

        # Loop through faces of m2.
        for u in m2.adj.keys():
            found = False
            vs = tg.get_edges(u)
            for v in vs:
                if (u, v) != (f1, f2):
                    if v in m1.adj:
                        old_adj = deepcopy(tf.adj)
                        tf.adj[f1].remove(f2)
                        tf.adj[f2].remove(f1)
                        tf.adj[u].add(v)
                        tf.adj[v].add(u)
                        refresh_mesh(tf)
                        tf.mk_xy_set()
                        if len(tf.xy_set) == 24:
                            print("Found a new mesh by connecting"
                                  + u + "," + v)
                            found = True
                            break
                        else:
                            print("Failed to find a valid mesh by connecting:"
                                  + u + "," + v)
                            tf.adj = old_adj
                            refresh_mesh(tf)
            if found:
                break

        if len(tf.xy_set) != 24:
            print("Invariant violated!")
            return
        if found:
            f_name = "Data//Meshes//mesh" +\
                                str(ii+base_i+1).rjust(9, '0') + ".txt"
            ii += 1
            adj = {}
            for k in tf.adj.keys():
                adj[k] = list(tf.adj[k])
            with open(f_name, "w") as file:
                file.write(json.dumps(adj))


def dedup_meshes(nu_strt_ix=0):
    """
    nu_strt_ix has to be 1 more than the number of distinct mesh
    files already in the directory. This will be printed before
    new meshes are generated by the method. Anything less than that
    will also work.
    """
    msh_files = os.listdir("Data/Meshes")
    msh_files = [i for i in msh_files if i.startswith("mesh")]
    msh_files = sorted(msh_files)
    to_del = []
    same_ones = defaultdict(set)
    for i in range(len(msh_files)):
        m1 = get_msh(msh_files[i])
        for j in range(max(nu_strt_ix,i+1), len(msh_files)):
            m2 = get_msh(msh_files[j])
            if m1.equals(m2):
                print("Will delete: " + msh_files[j])
                to_del.append(msh_files[j])
                same_ones[msh_files[j]].add(msh_files[i])
    for ff in set(to_del):
        f_name = "Data//Meshes//" + ff
        os.remove(f_name)
    for k in same_ones.keys():
        i = 0
        for ff in same_ones[k]:
            if i>0:
                f_name = "Data//Meshes//" + ff
                try:
                    print("And now deleting " + ff)
                    os.remove(f_name)
                except:
                    pass
            i += 1


def get_msh(filename):
    mesh_f1 = "Data//Meshes//" +\
                                filename
    with open(mesh_f1, "r") as file:
        adj_str = file.read()
    adj = json.loads(adj_str)
    adj1 = {}
    for k in adj.keys():
        adj1[k] = set(adj[k])
    tf = tg.TsrctFcGraph(angle=np.pi/2, adj=adj1)
    tf.dfs_flatten2('00+-')
    tf.actually_flatten()
    m1 = sm.SqMesh(tf)
    return m1


def cnt_mshs():
    files = os.listdir("Data/Meshes")
    files = [i for i in files if i.startswith("mesh")]
    return len(files)


def plot_mesh(ii=90):
    mesh_f1 = "Data//Meshes//mesh" +\
                                str(ii).rjust(4, '0') + ".txt"
    with open(mesh_f1, "r") as file:
        adj_str = file.read()
    adj = json.loads(adj_str)

    tf = tg.TsrctFcGraph(angle=np.pi/2)
    tf.adj = adj
    tf.r = rotation(4, np.pi*17/60.0)
    ot.open_given_cube(tf)


def tst_msh_eq():
    tf = tg.TsrctFcGraph(angle=np.pi/2)
    m1 = sm.SqMesh(tf)
    m1.vertices = np.array([[0,0],[0,-2],[0,-4],[-2,-4],[2,-4],[2,-6]])
    m1.process_verts()
    m2 = sm.SqMesh(tf)
    m2.vertices = np.array([[0,0],[0,-2],[-2,-2],[-2,-4],[-2,-6],[-4,-2]])
    m2.process_verts()
    m3 = sm.SqMesh(tf)
    m3.vertices = np.array([[0,0],[0,-2],[0,-4],[2,-4],[-2,-2],[-4,-2]])
    m3.process_verts()
    m4 = sm.SqMesh(tf)
    m4.vertices = np.array([[-2,2],[-2,0],[0,0],[2,0],[-2,-2],[-4,-2]])
    m4.process_verts()
    print(m1.equals(m2))
    print(m2.equals(m1))
    print(m3.equals(m1))
    print(m1.equals(m3))
    print(m4.equals(m1))
    print(m1.equals(m4))
