import sisl
import sisl.viz
# This is just for convenience to retreive files
siesta_files = sisl._environ.get_environ_variable("SISL_FILES_TESTS") / "sisl" / "io" / "siesta"

"""
Explanation: BandsPlot
End of explanation
"""


bands_plot = sisl.get_sile( siesta_files / "SrTiO3.bands").plot()

"""
Explanation: Let's get a bands_plot from a .bands file
End of explanation
"""


bands_plot

"""
Explanation: and see what we've got:
End of explanation
"""


bands_plot.update_settings(Erange=[-10, 10])

"""
Explanation: Getting the bands that you want
By default, BandsPlot gives you the 15 bands below and above 0 eV (which is interpreted as the fermi level). 
There are two main ways to specify the bands that you want to display: Erange and bands_range.
As you may have guessed, Erange specifies the energy range that is displayed:
End of explanation
"""


bands_plot.update_settings(bands_range=[6, 15], Erange=None)

"""
Explanation: while with bands_range you can actually indicate the indices.
However, note that Erange has preference over bands_range, therefore you need to set Erange to None if you want the change to take effect.
End of explanation
"""


bands_plot.update_settings(E0=-10, bands_range=None, Erange=None)

"""
Explanation: If your fermi level is not correctly set or you want a different energy reference, you can provide a value for E0 to specify where your 0 should be and the bands to display will be automatically calculated from that. 
However, if you want to update E0 after the plot has been build and you want BandsPlot to recalculate the bands for you you will need to set Erange and bands_range to None again.
End of explanation
"""


# Set them back to "normal"
bands_plot = bands_plot.update_settings(E0=0, bands_range=None, Erange=None)

"""
Explanation: Notice how only 25 bands are displayed now: the only 10 that are below 0 eV (there are no lower states) and 15 above 0 eV.
End of explanation
"""


bands_plot.update_settings(bands_color="red")

"""
Explanation: Notice that in spin polarized bands, you can select the spins to display using the spin setting, just pass a list of spin components (e.g. spin=[0]).
Quick styling
If all you want is to change the color and width of the bands, there's one simple solution: use the bands_color and bands_width settings.
Let's show them in red:
End of explanation
"""


bands_plot.update_settings(bands_color="green", bands_width=3)

"""
Explanation: And now in green but also make them wider:
End of explanation
"""


bands_plot = bands_plot.update_settings(bands_color="black", bands_width=1)

"""
Explanation: If you have spin polarized bands, bands_color will modify the color of the first spin channel, while the second one can be tuned with spindown_color.
End of explanation
"""


bands_plot.update_settings(gap=True, gap_color="green", Erange=[-10,10]) # We reduce Erange just to see it better

"""
Explanation: Displaying the smallest gaps
The easiest thing to do is to let BandsPlot discover where the (minimum) gaps are.
This is indicated by setting the gap parameter to True. One can also use gap_color if a particular color is desired.
End of explanation
"""


bands_plot.update_settings(direct_gaps_only=True)

"""
Explanation: This displays the minimum gaps. However there may be some issues with it: it will show all gaps with the minimum value. That is, if you have repeated points in the brillouin zone it will display multiple gaps that are equivalent. 
What's worse, if the region where your gap is is very flat, two consecutive points might have the same energy. Multiple gaps will be displayed one glued to another.
To help cope with this issues, you have the direct_gaps_only and gap_tol.
In this case, since we have no direct gaps, setting direct_gaps_only will hide them all:
End of explanation
"""


bands_plot.update_settings(direct_gaps_only=False, gap_tol=3)

"""
Explanation: This example is not meaningful for gap_tol, but it is illustrative of what gap_tol does. It is the minimum k-distance between two points to consider them "the same point" in the sense that only one of them will be used to show the gap. In this case, if we set gap_tol all the way up to 3, the plot will consider the two gamma points to be part of the same "point" and therefore it will only show the gap once.
End of explanation
"""


bands_plot = bands_plot.update_settings(gap=False, gap_tol=0.01)

"""
Explanation: This is not what gap_tol is meant for, since it is thought to remediate the effect of locally flat bands, but still you can get the idea of what it does.
End of explanation
"""


print(bands_plot.get_param("custom_gaps").help)

"""
Explanation: Displaying custom gaps
If you are not happy with the gaps that the plot is displaying for you or you simply want gaps that are not the smallest ones, you can always use custom_gaps.
Custom gaps should be a list where each item specifies how to draw that given gap. See the setting's help message:
End of explanation
"""


bands_plot.update_settings(custom_gaps=[{"from": "Gamma", "to": "Gamma", "color": "red"}])

"""
Explanation: So, for example, if we want to plot the gamma-gamma gap:
End of explanation
"""


bands_plot.bands.attrs

"""
Explanation: Notice how we got the gap probably not where we wanted, since it would be better to have it in the middle Gamma point, which is more visible. As the help message of custom_gaps states, you can also pass the K value instead of a label.
Now, you'll be happy to know that you can easily access the k values of all labels, as they are stored as attributes in the bands dataarray, which you can find in bands_plot.bands:
End of explanation
"""


gap_k = None
for val, label in zip(bands_plot.bands.attrs["ticks"], bands_plot.bands.attrs["ticklabels"]):
    if label == "Gamma":
        gap_k = val
gap_k

"""
Explanation: Now all we need to do is to grab the value for the second gamma point:
End of explanation
"""


bands_plot.update_settings(custom_gaps=[{"from": gap_k, "to": gap_k, "color": "orange"}])

"""
Explanation: And use it to build a custom gap:
End of explanation
"""


print(bands_plot.get_param("add_band_data").help)

"""
Explanation: Individual band styling
The bands_color and bands_width should be enough for most uses. However, you may want to style each band differently. Since we can not support every possible case, you can pass a function to the add_band_data. Here's the help message:
End of explanation
"""


def add_band_data(band, self):
    """Dummy function to see the band DataArray"""
    if band.band == 11:
        print(band)
    
    return {}

bands_plot.update_settings(add_band_data=add_band_data)

"""
Explanation: You can build a dummy function to print the band and see how it looks like. Notice that you only get those bands that are inside the range specified for the plot, therefore the first band here is band 11!
End of explanation
"""


import numpy as np
def draw_gradient(band, self):
    """
    Takes a band and styles it according to its energy dispersion.
    
    NOTE: If it's to far from the fermi level, it fades it in purple for additional coolness. 
    """
    dist_from_Ef = np.max(abs(band))
    if dist_from_Ef < 5:
        return {
            "mode": "lines+markers",
            "marker_size": np.abs(np.gradient(band))*40,
        }
    else:
        return {
            "line_color": "purple",
            "line_dash": "dot",
            "opacity": 1-float(dist_from_Ef/10)
        }
    
bands_plot.update_settings(add_band_data=draw_gradient)

"""
Explanation: Just as an educational example, we are going to style the bands according to this conditions:
- If the band is +- 5 eV within the fermi level, we are going to draw markers whose size is proportional to the gradient of the band at each point.
- Otherwise, we will just display the bands as purple dotted lines that fade as we get far from the fermi level (just because we can!)
Note: Of course, to modify traces, one must have some notion of how plotly traces work. Just hit plotly's visual reference page https://plotly.com/python/ for inspiration.
End of explanation
"""


H = sisl.get_sile(siesta_files / "Bi2D_BHex.TSHS").read_hamiltonian()
H.spin.is_spinorbit

"""
Explanation: Displaying spin texture
If your bands plot comes from a non-colinear spin calculation (or is using a Hamiltonian with non-colinear spin), you can pass "x", "y" or "z" to the spin setting in order to get a display of the spin texture.
Let's read in a hamiltonian coming from a spin orbit SIESTA calculation, which is obtained from this fantastic spin texture tutorial:
End of explanation
"""


band_struct = sisl.BandStructure(H, point=[[1./2, 0., 0.], [0., 0., 0.],
                                     [1./3, 1./3, 0.], [1./2, 0., 0.]],
                           division=301,
                           name=['M', r'$\Gamma$', 'K', 'M'])

"""
Explanation: Generate the path for our band structure:
End of explanation
"""


spin_texture_plot = band_struct.plot(Erange=[-2,2])
spin_texture_plot

"""
Explanation: And finally generate the plot:
End of explanation
"""


spin_texture_plot.update_settings(spin="x", bands_width=3)

"""
Explanation: These are the bands, now let's ask for a particular spin texture:
End of explanation
"""


spin_texture_plot.update_settings(backend="plotly", spin_texture_colorscale="temps")

"""
Explanation: And let's change the colorscale for the spin texture:
End of explanation
"""


thumbnail_plot = spin_texture_plot

if thumbnail_plot:
    thumbnail_plot.show("png")

"""
Explanation: We hope you enjoyed what you learned!

This next cell is just to create the thumbnail for the notebook in the docs
End of explanation
"""


