import numpy as np
np.float = float

import defdap.hrdic as hrdic
import defdap.ebsd as ebsd
from defdap.quat import Quat

%matplotlib inline

ebsd_map = ebsd.Map('../tests/data/testDataEBSD')
ebsd_map.buildQuatArray()
ebsd_map.findBoundaries(boundDef=8)
ebsd_map.findGrains(minGrainSize=10)
ebsd_map.calcGrainAvOris()
ebsd_map.buildNeighbourNetwork()
ebsd_map.homogPoints = [
    (68, 95),
    (308, 45),
    (191, 187),
    (89, 174)
]

dic_map = hrdic.Map('../tests/data/', 'testDataDIC.txt')
dic_map.setScale(20 / 2048)
dic_map.homogPoints = [
    (36, 72),
    (279, 27),
    (162, 174),
    (60, 157)
]
dic_map.linkEbsdMap(ebsd_map, transformType="affine")
dic_map.findGrains(minGrainSize=10)

# find twin boundaries
misori_twin = Quat.fromAxisAngle([1, 1, 1], 60*np.pi/180)
misori_twin_tol = 5*np.pi/180

# create all symmetric equivalent misorientations
misori_twin_all = []
syms = ebsd_map.primaryPhase.crystalStructure.symmetries
for sym_i in syms:
    for sym_j in syms:
        misori_twin_all.append(sym_i.conjugate * misori_twin * sym_j)
# get rid of any duplicates
misori_twin_all = list(set(misori_twin_all))

# loop over all grain boundary segments and check if the misorientation between
# the two grains is within tolerance of the twin misorientation
twin_lines = []
for grain1, grain2, b_seg in ebsd_map.neighbourNetwork.edges.data('boundary'):
    twin = False
    misori = grain2.refOri * grain1.refOri.conjugate
    for misori_twin in misori_twin_all:
        if 2 * np.arccos(misori_twin.dot(misori)) < misori_twin_tol:
            twin = True
            break
    
    if not twin:
        continue
    
    twin_lines += b_seg.boundaryLines

plot = ebsd_map.plotEulerMap(plotGBs='line')
plot.addGrainBoundaries(boundaries=twin_lines, kind='line', colour='limegreen')

plot = dic_map.plotMaxShear(vmin=0, vmax=0.05, plotGBs='line')
plot.addGrainBoundaries(
    boundaries=dic_map.warp_lines_to_dic_frame(twin_lines), 
    kind='line', colour='limegreen', lw=3
)

import matplotlib.pyplot as plt
import matplotlib as mpl
import defdap.plotting

def oris_to_col(ori_1: Quat, ori_2: Quat, sym_group: str) -> float:
    return ori_1.misOri(ori_2, sym_group)

sym_group = ebsd_map.primaryPhase.crystalStructure.name
gb_lines = []
gb_cols = []
for grain1, grain2, b_seg in ebsd_map.neighbourNetwork.edges.data('boundary'):
    line_col = oris_to_col(grain1.refOri, grain2.refOri, sym_group)
    gb_lines += b_seg.boundaryLines
    gb_cols += [line_col] * len(b_seg.boundaryLines)

# Convert values to colours
gb_cols = 2 * np.arccos(gb_cols)

plot = ebsd_map.plotEulerMap()
img = plot.addGrainBoundaries(
    kind='line',
    boundaries=gb_lines,
    colour=gb_cols,
    cmap=mpl.colormaps['gray'],
    norm=mpl.colors.Normalize(vmin=gb_cols.min(), vmax=gb_cols.max()),
)
plt.colorbar(mappable=img, label="Boundary misorientation")

plot = defdap.plotting.MapPlot.create(dic_map, np.zeros(dic_map.shape))
img = plot.addGrainBoundaries(
    kind='line',
    boundaries=dic_map.warp_lines_to_dic_frame(gb_lines),
    colour=gb_cols,
    cmap=mpl.colormaps['viridis'],
    norm=mpl.colors.Normalize(vmin=gb_cols.min(), vmax=gb_cols.max()),
)
plt.colorbar(mappable=img, label="Boundary misorientation")

plot = dic_map.plotMaxShear(vmin=0, vmax=0.05)
img = plot.addGrainBoundaries(
    kind='line',
    boundaries=dic_map.warp_lines_to_dic_frame(gb_lines),
    colour=gb_cols,
    cmap=mpl.colormaps['gray'],
    norm=mpl.colors.Normalize(vmin=gb_cols.min(), vmax=gb_cols.max()),
)
plt.colorbar(mappable=img, label="Boundary misorientation")