# Import potrzebnych modułów
%matplotlib inline
import numpy as np
from ase import Atoms, units
import ase.io
from ase.io.trajectory import Trajectory
from ipywidgets import HBox, VBox, Checkbox, Dropdown, IntSlider, FloatSlider
from io import BytesIO
import nglview 
import glob

def recenter(a):
    '''
    Normalizacja położenia nanocząstki do stałej pozycji środka masy.
    Przemieszczenie centrum masy do środka komórki.
    Uwaga: do działania konieczna jest istotna próżnia wokoło struktury.
    '''
    # Kopia struktury a
    c=Atoms(numbers=a.get_atomic_numbers(),
            positions=a.get_positions(),
            cell=a.get_cell(),
            pbc=a.get_pbc())
    
    c.translate((c.get_cell()/2).sum(axis=0)-c.get_center_of_mass())
    c.set_scaled_positions(c.get_scaled_positions())
    c.translate((c.get_cell()/2).sum(axis=0)-c.get_center_of_mass())
    c.set_scaled_positions(c.get_scaled_positions())
    return c

"""
Explanation: Warsztaty modelowania w nanofizyce

Zachowania atomów w zależności od ich rodzaju i położenia
Paweł T. Jochym
Zakład Komputerowych Badań Materiałów
Instytut Fizyki Jądrowej PAN, Kraków

Analiza przeprowadzona w ćwiczeniu Amplitudy wskazuje, że atomy w nanocząstce Pt-12Fe-42Pt zachowują się bardzo różnie w zależności od tego czy są to atomy żelaza czy platyny. Postarajmy się zbadać jak wygląda to zachowanie rysując struktury selekcjonując atomy jednego rodzaju. Zadanie selekcji wykona funkcja wyświetlania struktury.
Poniższy zeszyt demonstruje proste i bardziej złożone operacje jakie można zrealizować przy pomocy odpowiedniej wizualizacji. Kod zawarty poniżej można wykorzystać jako przykład do dalszej modyfikacji pozwalający przedstawić potrzebne cechy badanego układu.
End of explanation
"""


mdtraj={}
print('Czytanie trajektorii:', end=' ')

for fn in glob.glob('data/md_T_*.traj'):
    # Użyj części nazwy pliku jako identyfikacji temperatury
    T=int(fn.split('/')[-1][5:-5])
    # Wczytaj trajektorię z pliku
    print(T, end=' ')
    mdtraj[T]=[recenter(a) for a in Trajectory(fn)]

print()

"""
Explanation: Wczytanie wyliczonych trajektorii z dysku z jednoczesnym wycentrowaniem nanocząstek w komórce elementarnej.
End of explanation
"""


def show_cryst(struct, uc=True, re_center=False):
    '''
    Proste wyświetlenie struktury krystalicznej. 
    Opcjonalnie wraz z komórką elementarną i re-centrowaniem nanocząstki.
    '''
    view = nglview.show_ase(recenter(struct) if re_center else struct)
    view.parameters=dict(clipDist=-100)
    if uc : view.add_unitcell()
    view.camera='orthographic'
    view.add_spacefill(radiusType='covalent', scale=0.5)
    view.center()
    return view  

class TrajDisplay:
    '''
    Wyświetlacz trajektorii z możliwością selekcji atomów 
    w.g. pierwiastka oraz trajektorii z całego zbioru przekazanego
    jako słownik. Opcjonalnie całość trajektorii może zostać wycentrowana
    w komórce periodycznej.
    '''
    def __init__(self, trajectories, re_center=False):
        if re_center :
            print('Re-centering...')
            self.trajectories = {k:[recenter(a) for a in trj] 
                                     for k, trj in trajectories.items()}
        else :
            self.trajectories=trajectories
        self.temperatures=sorted(self.trajectories.keys())
        self.T=self.temperatures[0]
        self.trj=self.trajectories[self.T]

        self.v=nglview.show_asetraj(self.trj)
        self.v._remote_call("setSize", target="Widget", args=["500px", "500px"])
        self.v.add_spacefill()
        self.v.parameters=dict(clipDist=-100)
        self.v.camera='orthographic'
        self.v.update_spacefill(radiusType='covalent', scale=0.7)
        self.v.center()
        
        self.tsel=Dropdown(options=self.temperatures,
                           value=self.T, 
                           description='Temperatura:')
        self.frm=IntSlider(value=0, min=0, max=len(self.trj)-1)
        
        self.asel=Dropdown(options=['All']+list(set(self.trj[0].get_chemical_symbols())),
                           value='All', description='Wyświetl')
        
        self.rad=FloatSlider(value=0.8, min=0.0, max=1.5, step=0.01, description='Rozmiar')

        self.tsel.observe(self._update_trj, 'value')
        self.frm.observe(self._update_frame)
        self.asel.observe(self._select_atom)
        self.rad.observe(self._update_repr)
        self._update_trj(None)
        
        self.hbox = HBox([self.v, VBox([self.tsel, 
                                        self.asel, 
                                        self.rad,
                                        self.frm])])

        
    def _update_repr(self, chg=None):
        self.v.update_spacefill(radiusType='covalent', scale=self.rad.value)        
        
    def _update_trj(self, chg=None):
        '''Zmiana trajektorii zależnie od wybranej temperatury'''
        self.T=int(self.tsel.value)
        self.trj=self.trajectories[self.T]
        self.frm.max = len(self.trj)-1
        v=self.v
        while v._ngl_component_ids :
            v.remove_component(v._ngl_component_ids[0])
        v.add_trajectory(nglview.ASETrajectory(self.trj))
        v.add_unitcell()
        v.center()
        self._select_atom()
        self._update_frame()
        return
    

    def _update_frame(self, chg=None):
        self.v.frame=self.frm.value
        return
    
    
    def _select_atom(self, chg=None):
        sel=self.asel.value
        self.v.remove_spacefill()
        if sel == 'All':
            self.v.add_spacefill(selection='All')
        else :
            self.v.add_spacefill(selection=[n for n,e in enumerate(self.trj[0].get_chemical_symbols()) if e == sel])
        self._update_repr()
            


"""
Explanation: Funkcje wizualizujące. Proste wyświetlanie struktury oraz bardziej złożony wizualizator trajektorii z możliwością wyboru rodzaju pokazywanych atomów oraz wyświetlanego kroku symulacji.
End of explanation
"""


show_cryst(mdtraj[1400][700])

dsp=TrajDisplay(mdtraj)
dsp.hbox

"""
Explanation: Wyświetlenie struktur przy użyciu zdefiniowanych powyżej narzędzi.
End of explanation
"""


