from ipyparallel import Client
cluster = Client()
dview = cluster[:]
dview.use_dill()
lview = cluster.load_balanced_view()
len(dview)

"""
Explanation: Phase diagram for multiple angles
Start a ipcluster from the Cluster tab in Jupyter or use the command:
ipcluster start 
in a terminal.
End of explanation
"""


# import os
# from scripts.hpc05 import HPC05Client
# os.environ['SSH_AUTH_SOCK'] = os.path.join(os.path.expanduser('~'), 'ssh-agent.socket')
# cluster = HPC05Client()

"""
Explanation: This next cell is for internal use with our cluster at the department, a local ipcluster will work: use the cell above.
End of explanation
"""


%%px --local
import sys
import os
# CHANGE THE LINE BELOW INTO THE CORRECT FOLDER!
sys.path.append(os.path.join(os.path.expanduser('~'), 'orbitalfield'))
import kwant
import numpy as np
from fun import *

def gap_and_decay(lead, p, val, tol=1e-4):
    gap = find_gap(lead, p, val, tol)
    decay_length = find_decay_length(lead, p, val)
    return gap, decay_length

import holoviews as hv
import holoviews_rc
hv.notebook_extension()

"""
Explanation: Make sure to add the correct path like:
sys.path.append("/path/where/to/ipynb/runs")
End of explanation
"""


%%px --local
# angle = 0 # WIRE WITH SC ON TOP

angle = 45 # WIRE WITH SC ON SIDE
p = make_params(t_interface=7/8*constants.t, Delta=68.4, r1=50, r2=70, 
                orbital=True, angle=angle, A_correction=True, alpha=100) #r2=70

p.V = lambda x, y, z: 2 / 50 * z
lead = make_3d_wire_external_sc(a=constants.a, r1=p.r1, r2=p.r2, angle=p.angle)

# WIRE WITH CONSTANT GAP
# lead = make_3d_wire()
# p = make_params(V=lambda x, y, z: 0, orbital=True)

"""
Explanation: Uncomment the lines for the wire that you want to use.
End of explanation
"""


# give an array of angles that you want to use

# thetas = np.array([0, np.tan(1/10), 0.5 * np.pi - np.tan(1/10), 0.5 * np.pi])
# phis = np.array([0, np.tan(1/10), 0.5 * np.pi - np.tan(1/10), 0.5 * np.pi])

thetas = np.array([0.5 * np.pi])
phis = np.array([0])

# the range of magnetic field and chemical potential
Bs = np.linspace(0, 2, 400)
mu_mesh = np.linspace(0, 35, 400)

# creates a 3D array with all values of magnetic field for all specified angles
pos = spherical_coords(Bs.reshape(-1, 1, 1), thetas.reshape(1, -1, 1), phis.reshape(1, 1, -1))
pos_vec = pos.reshape(-1, 3)

mus_output = lview.map_sync(lambda B: find_phase_bounds(lead, p, B, num_bands=40), pos_vec)
mus, vals, mask = create_mask(Bs, thetas, phis, mu_mesh, mus_output)

N = len(vals)
step = N // (len(phis) * len(thetas))
print(N, step)

"""
Explanation: You can specify the angles that you want to calculate in thetas and phis.
Also specify the range of magnetic field and chemical potential in Bs and mu_mesh.
End of explanation
"""


import holoviews_rc
from itertools import product
from math import pi

kwargs = {'kdims': [dimensions.B, dimensions.mu],
          'extents': bnds(Bs, mu_mesh),
          'label': 'Topological boundaries',
          'group': 'Lines'}

angles = list(product(enumerate(phis), enumerate(thetas)))

boundaries = {(theta / pi, phi / pi): hv.Path((Bs, mus[i, j, :, ::2]), **kwargs)
                  for (i, phi), (j, theta) in angles}

BlochSpherePlot.bgcolor = 'white'

sphere = {(theta / pi, phi / pi): BlochSphere([[1, 0, 0], spherical_coords(1, theta, phi)], group='Sphere')
              for (i, phi), (j, theta) in angles}

hv.HoloMap(boundaries, **dimensions.angles) + hv.HoloMap(sphere, **dimensions.angles)

"""
Explanation: Check whether the correct angles were used and see the phase boundaries
End of explanation
"""


full_phase_diagram = False

"""
Explanation: Calculate full phase diagram
Make sure tempdata exists in the current folder. 
Set full_phase_diagram to False if you only want the band gap in the non-trivial region or True if you want it in the whole Bs, mu_mesh range.
End of explanation
"""


import os.path
import sys

fname_list = []
for i, n in enumerate(range(0, N, step)):
    fname = "tempdata/" + str(n)+"-"+str((i+1)*step)+".dat"
    fname_list.append(fname)
    
    if not os.path.isfile(fname):  # check if file already exists
        lview.results.clear()
        cluster.results.clear()
        cluster.metadata.clear()
        print(fname)
        sys.stdout.flush()
        if full_phase_diagram:
            gaps_and_decays_output = lview.map_async(lambda val: gap_and_decay(lead, p, val[:-1] + (True,)), vals[n:(i+1) * step])
        else:
            gaps_and_decays_output = lview.map_async(lambda val: gap_and_decay(lead, p, val), vals[n:(i+1) * step])
        gaps_and_decays_output.wait_interactive()
        np.savetxt(fname, gaps_and_decays_output.result())
        print(n, (i+1) * step)
cluster.shutdown(hub=True)

gaps_and_decay_output = np.vstack([np.loadtxt(fname) for fname in fname_list])
gaps_output, decay_length_output = np.array(gaps_and_decay_output).T

gaps = np.array(gaps_output).reshape(mask.shape)
gaps[1:, 0] = gaps[0, 0]

decay_lengths = np.array(decay_length_output).reshape(mask.shape)
decay_lengths[1:, 0] = decay_lengths[0, 0]

if full_phase_diagram:
    gaps = gaps*(mask*2 - 1)
    decay_lengths = decay_lengths*(mask*2 - 1)
    gaps_output = gaps.reshape(-1)
    decay_length_output = decay_lengths.reshape(-1)

"""
Explanation: The next cell calculates the gaps and decay lengths.
You can stop and rerun the code, it will skip over the files that already exist.
Make sure the folder tempdata/ exists.
End of explanation
"""


fname = 'data/test.h5'
save_data(fname, Bs, thetas, phis, mu_mesh, mus_output, gaps_output, decay_length_output, p, constants)

"""
Explanation: Save
Run this function to save the data to hdf5 format, it will include all data and parameters that are used in the simulation.
End of explanation
"""


%%output size=200
%%opts Image [colorbar=False] {+axiswise} (clims=(0, 0.1))
phase_diagram = create_holoviews(fname)

(phase_diagram.Phase_diagram.Band_gap.hist()
 + phase_diagram.Phase_diagram.Inverse_decay_length 
 + phase_diagram.Sphere.I).cols(2)

%%opts Image [colorbar=True]
phase_diagram.Phase_diagram.Band_gap

phase_diagram.cdims

"""
Explanation: Check how the phase diagram looks
This will show all data.
End of explanation
"""


