#https://qiita.com/ojiya/items/b97d25b896741ad99b79
#https://qiita.com/ojiya/items/a1f1dd358de1231ea157
# https://qiita.com/ojiya/items/9ba8bfadab96cdddb0b9
#Pymatgenチュートリアル④ 構造を扱う https://qiita.com/ojiya/items/e98a9dd5cb6cd7ad38ca
#Pymatgenチュートリアル⑤ Structure型のデータを作る https://qiita.com/ojiya/items/37b594150115531481f0
#Pymatgenチュートリアル⑥ XRDのシミュレーションをする https://qiita.com/ojiya/items/1b154c3698cff91c8a2b


import math
import numpy as np

from pymatgen.core.periodic_table import Element
from pymatgen.core.composition import Composition
from pymatgen.io.cif import CifParser
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.core.composition import Composition


lat = Lattice.from_parameters(4,5,6,90,90,90)
#lat_cubic = Lattice.cubic(8.53335154)
#lat = Lattice.orthorhombic(4,5,6)
print(type(lat))


path = 'SrTiO3.cif'
parser = CifParser(path)
mat = parser.get_structures(primitive=False, symmetrized=False)[0]
lat_SrTiO3 = mat.lattice
print(lat_SrTiO3)

sga = SpacegroupAnalyzer(mat)
mat_sym = sga.get_symmetrized_structure()
eq_sites = mat_sym.equivalent_sites

asym_unit = [site[0] for site in eq_sites]


#対称操作を取り出す
sym_ops = sga.get_space_group_operations()
op = sym_ops[1]
rot = op.rotation_matrix
trans = op.translation_vector

# もとになるサイトを取り出す。
site_0 = asym_unit[0]
orig_coord = site_0.frac_coords
print(orig_coord)

# 対称操作を行う
gen_sites = []
for op in sym_ops:
    rot = op.rotation_matrix
    trans = op.translation_vector
    gen_coord = [round(coord-math.floor(coord),4) for coord in np.dot(rot, orig_coord) + trans]
    gen_sites.append(gen_coord)

print(gen_sites[0])
print(gen_sites[21])


gen_sites = []
for op in sym_ops:
    rot = op.rotation_matrix
    trans = op.translation_vector
    gen_coord = [coord-math.floor(coord) for coord in np.dot(rot, orig_coord) + trans]
    gen_sites.append(gen_coord)

print(gen_sites[0])
print(gen_sites[21])


def get_unique_list(seq):
    seen = []
    return [x for x in seq if x not in seen and not seen.append(x)]
    
def gen_sym_sites(orig_coord, set_sym_ops):
    gen_sites = []
    for op in sym_ops:
        rot = op.rotation_matrix
        trans = op.translation_vector
        gen_coord = [round(coord-math.floor(coord),4) for coord in np.dot(rot, orig_coord) + trans]
        gen_sites.append(gen_coord)
    unique_sites = get_unique_list(gen_sites)
    return unique_sites    


gen_sites = gen_sym_sites(orig_coord, sym_ops)

dict_gen_sites = {}
for i, site in enumerate(asym_unit):
    orig_coord = site.frac_coords
    gen_sites = gen_sym_sites(orig_coord, sym_ops)
    gen_key = f'{str(site.species)}_{i}'
    dict_gen_sites[gen_key] = gen_sites
dict_gen_sites


key_list = list(dict_gen_sites.keys())
list_species = []
list_coord = []
for k in key_list:
    spec = Composition(k.split('_')[0])
    list_gen_sites = dict_gen_sites[k]
    n_gen_site =len(list_gen_sites)
    list_species += [spec] * n_gen_site
    list_coord += list_gen_sites
    
from pymatgen.core.structure import Structure
gen_structure = Structure(lat, list_species, list_coord)
print(gen_structure)
    
    
print("END")

    
    


"""
path = 'SrTiO3.cif'
parser = CifParser(path)
mat = parser.get_structures(primitive=False, symmetrized=False)[0]
print(type(mat))
dict_mat = mat.as_dict()
print(dict_mat.keys())
print(dict_mat)

mat_lattice = mat.lattice
print(type(mat_lattice))
print(mat_lattice)
print(mat_lattice.parameters)
print(mat_lattice.a)
print(mat_lattice.alpha)

mat_spec = mat.species
print(mat_spec)
mat_spec_occ = mat.species_and_occu
print(mat_spec_occ)

mat_sites = mat.sites
print(type(mat_sites))
print(type(mat_sites[0]))
print(mat_sites[0])
print(mat_sites[0].species)
print(mat_sites[0].coords)
print(mat_sites[0].frac_coords)


mat_sg = mat.get_space_group_info()
print(mat_sg)
sga = SpacegroupAnalyzer(mat)
print(sga)
mat_sym = sga.get_symmetrized_structure()
eq_sites = mat_sym.equivalent_sites

print(type(eq_sites))
print(type(eq_sites[0]))
print(eq_sites[0])
asym_unit = [site[0] for site in eq_sites]
print(asym_unit)

sym_ops = sga.get_symmetry_operations()
print(type(sym_ops))
print(type(sym_ops[0]))
print(sym_ops[1].rotation_matrix)
print(sym_ops[1].translation_vector)
op = sym_ops[1]
rot = op.rotation_matrix
trans = op.translation_vector
print(rot)
print(trans)

frac_coord = asym_unit[0].frac_coords
print(frac_coord)

sym_generated_raw = np.dot(rot, frac_coord) + trans
print(sym_generated_raw)

sym_generated = [coord - math.floor(coord) for coord in sym_generated_raw]
print(sym_generated)
"""



"""
mat0 = parser.get_structures()[0]
print("")
print(type(mat0))
print(mat0.as_dict())

print("")
print(mat0)

print(mat0.formula)
print(mat0.get_space_group_info())
print(mat0.lattice)

print("")
print("Charge guess")
ox_guess = mat0.composition.oxi_state_guesses()
print("Guessed charges: ", ox_guess)
print("Total charge:", mat0.charge)

print("")
print("Charge specify")
charges = {'Sr': 3.0, 'Ti': 3.0, 'O': -1.0}
mat0.add_oxidation_state_by_element(charges)
print("Specified charges: ", charges)
print("Total charge:", mat0.charge)

print("")
print(type(mat0.sites))
print(mat0.sites)
print(mat0.sites[0])

metal_sites = []
for site in mat0.sites:
    if site.specie.is_metal:
       metal_sites.append(site) 
print(metal_sites)

#distance：2つのサイト間の距離を求める。
#round(少数点数、整数)：少数点数を、整数桁目で四捨五入した値を得る。
mm_dist = [round(metal1.distance(neighbor), 4) for neighbor in neighbors_metal]

#numpyのnonzeroを使って、0以外の値を取り出す。
mm_dist_nonzero = np.array(mm_dist)[np.nonzero(mm_dist)]

#minを使って、最小値を得る。
mm_min = min(mm_dist_nonzero)
print(mm_min)

# 内部座標を還元
print("")
metal1 = metal_sites[0]
neighbors = mat0.get_neighbors(metal1, 5)
print(neighbors)
"""






"""
#普通の元素記号でデータを持ってくることができる
print("")
print("Element: Si")
si = Element('Si')
print(type(si))
print(type(si.data))
print(si.data)
print(si.full_electronic_structure)
print(si.is_metal)
print(si._atomic_radius)


#Compositionの中に、文字列で普通に組成式を入れればいい
print("")
print("Compound: Ti3O5")
ti3o5 = Composition("Ti3O5")
print(type(ti3o5))
print(ti3o5.to_data_dict)
print(ti3o5.elements)
print(ti3o5.to_reduced_dict)
print(ti3o5.to_reduced_dict.keys())
print(ti3o5.reduced_formula)


#よくある小数点の添字をもった化合物を扱う
#半端な組成の化合物を持ってくる
mat_3_0 = Composition('Ti2.98O5.10')

#添字の数字だけを取り出す
sub_num = mat_3_0.to_reduced_dict.values()
#添字の数字の中の最大値で添字を全て割る = 規格化
max_sub = max(sub_num)
std_sub = [i/max_sub for i in sub_num]

#規格化した添字を持つ式に直す
#もうちょいうまい書き方があるが、ここはこれでいいことにする
mat_3_1 = Composition(str(mat_3_0.elements[0]) + str(std_sub[0])
+str(mat_3_0.elements[1]) + str(std_sub[1]))
print(mat_3_1)

#モデルにしたい化合物の組成式を持ってきて、上と同じことをする
mat_4_0 = Composition('Ti3O5')
sub_num = mat_4_0.to_reduced_dict.values()
max_sub = max(sub_num)
std_sub = [i/max_sub for i in sub_num]

mat_4_1 = Composition(str(mat_4_0.elements[0]) + str(std_sub[0])
+str(mat_4_0.elements[1]) + str(std_sub[1]))
print(mat_4_1)

#2つの式が、ある許容範囲以内で同じものかを判別する。
#判断したい式.almost_equals(基準にする式)という式で、bool型の値を返す
#例えば、論文に登場した式が、あるデータベースに含まれるか判断したいときに使える
print(mat_4_1.almost_equals(mat_4_1))

#出力結果：Ti0.58431373 O1
#Ti0.6 O1
#True
"""