import json

import pytest
from dbt.tests.util import run_dbt

ref_models__table_comment_sql = """
{{
  config(
    materialized = "table",
    persist_docs = {"relation": true, "columns": true},
  )
}}

select
    'foo' as first_name,
    'bar' as second_name

"""

ref_models__view_comment_sql = """
{{
  config(
    materialized = "view",
    persist_docs = {"relation": true, "columns": true},
  )
}}

select
    'foo' as first_name,
    'bar' as second_name

"""

ref_models__schema_yml = """
version: 2

models:
  - name: table_comment
    description: "YYY table"
    columns:
      - name: first_name
        description: "XXX first description"
      - name: second_name
        description: "XXX second description"
  - name: view_comment
    description: "YYY view"
    columns:
      - name: first_name
        description: "XXX first description"
      - name: second_name
        description: "XXX second description"
"""

class TestBaseComment:
    @pytest.fixture(scope="class")
    def models(self):
        return {
            "schema.yml": ref_models__schema_yml,
            "table_comment.sql": ref_models__table_comment_sql,
            "view_comment.sql": ref_models__view_comment_sql,
        }

    @pytest.mark.parametrize(
        'model_name',
        ['table_comment', 'view_comment'],
    )
    def test_comment(self, project, model_name):
        run_dbt(["run"])
        run_dbt(["docs", "generate"])
        with open("target/catalog.json") as fp:
            catalog_data = json.load(fp)

        assert "nodes" in catalog_data
        column_node = catalog_data["nodes"][f"model.test.{model_name}"]
        for column in column_node["columns"].keys():
            column_comment = column_node["columns"][column]["comment"]
            assert column_comment.startswith("XXX")

        assert column_node['metadata']['comment'].startswith("YYY")


import numpy as np
import matplotlib.pyplot as plt
from scipy.io import wavfile

def plot_fourier_spectrum(freq, spectrum, title):
    plt.plot(freq, np.abs(spectrum))
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Magnitude")
    plt.title(title)

# Load the WAV file
sample_rate, audio_data = wavfile.read("C:/Users/barwa/Desktop/Math Project/InputData.wav.wav")

# Extract the mono audio channel (if stereo)
if len(audio_data.shape) > 1:
    audio_data = audio_data[:, 0]

# Perform the DFT
spectrum = np.fft.fft(audio_data)
freq = np.fft.fftfreq(len(audio_data), d=1/sample_rate)

# Reconstruct the signal using only the first five Fourier coefficients
reconstructed_spectrum = np.zeros_like(spectrum)
reconstructed_spectrum[:5] = spectrum[:5]
reconstructed_signals = []
for i in range(5):
    reconstructed_spectrum_i = np.zeros_like(spectrum)
    reconstructed_spectrum_i[i] = spectrum[i]
    reconstructed_signal_i = np.fft.ifft(reconstructed_spectrum_i)
    reconstructed_signals.append(reconstructed_signal_i)

# Plot the original waveform
plt.figure(figsize=(8, 6))
plt.subplot(6, 1, 1)
plt.plot(audio_data)
plt.xlabel("Time (samples)")
plt.ylabel("Amplitude")
plt.title("Original Waveform")

# Plot the reconstructed waveforms
for i in range(5):
    plt.subplot(6, 1, i+2)
    plt.plot(reconstructed_signals[i].real)
    plt.xlabel("Time (samples)")
    plt.ylabel("Amplitude")
    plt.title("Reconstructed Waveform ({} Fourier coefficient)".format(i+1))

plt.tight_layout()
plt.show()

# Plot the Fourier spectra of the reconstructed waveforms
plt.figure(figsize=(8, 6))
for i in range(5):
    plt.subplot(5, 1, i+1)
    reconstructed_spectrum_i = np.zeros_like(spectrum)
    reconstructed_spectrum_i[i] = spectrum[i]
    reconstructed_signal_i = np.fft.ifft(reconstructed_spectrum_i)
    reconstructed_spectrum_i = np.fft.fft(reconstructed_signal_i)
    plot_fourier_spectrum(freq, reconstructed_spectrum_i, "Frequency Spectrum ({} Fourier coefficient)".format(i+1))

plt.tight_layout()
plt.show()


#  メッシュコード計算

#  meshcode_to_latlong,  メッシュコードから緯度・経度へ
#  latlong_to_meshcode,  緯度・経度からメッシュコードへ。
#  dms_to_deg,           度・分・秒から度（小数点表示）へ。
#  deg_to_dms,           度（小数点表示）から度・分・秒へ。

#  latlong_to_half_meshcode, 　緯度・経度から1/2メッシュコードを作成

import re  # 正規表現によるパターンマッチ

########################################
# メッシュコードを渡し、対応する緯度・経度を返す。
# メッシュコードは数値（整数）でも文字列型でも可。
#
# 返値は数値データが２つ並んだタプル。
#
# 特に指定しなければメッシュの中央の座標を返す。
# メッシュコードは一次、ニ次、三次いずれも可。
#
# NW, NE, SW, SE を２つ目の引数として渡すと、それぞれ北西、北東、南西、南東の
# コーナーの座標を返す。
#
# 緯度・経度は度単位の小数。
#
#  usage
#  code_a = "55354251" # 三次メッシュコード
#  lat_lon = mesh_lib.meshcode_to_latlong(code_a);  # 中央の緯度・経度
#  print('latitude\t', lat_lon[0], '\tlongitude\t', lat_lon[1],
#
#  code_a = "55354251" # 三次メッシュコード
#  lat_lon = mesh_lib.meshcode_to_latlong(code_a, 'NE') # 北東端
#  print('latitude\t', lat_lon[0], '\tlongitude\t', lat_lon[1],

def meshcode_to_latlong(code, loc = 'C'):

    if type(code) != int:  # 整数型に（文字列が渡された場合に対応）
        try:
            code = int(code)
        except:  # 形式チェック
            print("Error in  mesh_lib.meshocode_to_latlong()  Not a valid code handed.")
            raise

    code = str(code) # あらためて文字型に。

    code12 = ''
    code34 = ''
    code5  = ''
    code6  = ''
    code7  = ''
    code8  = ''

    loc = loc.upper()

    if re.match("^C|(NE)|(NW)|(SE)|(SW)$", loc) == None:
        raise Exception ("Invalid option to meshcode_to_latlong()")

    match_1_result = re.match("\d{4}", code)

    if re.match("\d{4}", code):  # 最初の４文字が数字

        code12 = code[0:2]
        code34 = code[2:4]
        lat_width  = 2.0 / 3.0  # grid cell の緯度方向の間隔
        long_width = 1.0        # grid cell の経度方向の間隔

    else:
        return (None)           # メッシュコードとして無効

    if re.match("\d{6}", code):   # 少なくとも最初の６文字は数字

        code5 = code[4:5]
        code6 = code[5:6]
        lat_width  /= 8.0;
        long_width /= 8.0;

    if re.match("\d{8}", code):  # 最初の８文字は数字

        code7 = code[6:7]
        code8 = code[7:8]
        lat_width  /= 10.0;
        long_width /= 10.0;

    # 以下、南西コーナーの座標を求める。

    lat  = float(code12) * 2 / 3          #  一次メッシュ
    long = float(code34) + 100

    if (code5 != '') & (code6 != ''):     # 二次メッシュ or 三次メッシュ
        lat  += (float(code5) * 2 / 3) / 8
        long += float(code6) / 8

    if (code7 != '') & (code8 != ''):     # 三次メッシュ
        lat  += float(code7) * 2 / 3 / 8 / 10
        long += float(code8) / 8 / 10

    # ここまでで南東端の緯度・経度が得られた

    # 中央の座標なら、一区画の幅（経度幅）・高さ（緯度幅）の半分を足す。
    if loc == 'C':
        lat  += lat_width  / 2
        long += long_width / 2

    # 北端の座標なら、一区画の高さ（緯度幅）を足す
    if re.search('N', loc):
        lat += lat_width

    # 東端の座標なら、一区画の幅（経度幅）を足す
    if re.search('E', loc):
        long  += long_width

    return (lat, long)     # タプルを返す。

##################
#  緯度・経度（度単位）を受けとり、これを含むメッシュのコードを返す
#  緯度。経度の型は float でも文字列でも可。
#
#  返値のメッシュコードは数値（整数）データ。
#
#  ３番めの引数として 1,2,3 を渡せば、一次メッシュ、二次メッシュ、
#  三次メッシュコードを計算する。指定がなければ三次メッシュ、
#
# Usage:
#  code3 = meshlib.latlong_to_meshcode(35.3, 138.5)      # 三次メッシュコード
#  code2 = meshlib.latlong_to_meshcode(35.3, 138.5, 2)   # 二次メッシュコード
#  code1 = meshlib.latlong_to_meshcode(35.3, 138.5, 1)   # 一次メッシュコード
#

def latlong_to_meshcode(lat, long, order = 3):

    if type(lat) != float:  # 実数型でないなら変換
        lat = float(lat)

    if type(long) != float: # 実数型でないなら変換
        long = float(long)

    if (order < 1) | (order > 3):
        raise Exception ("Invalid parameter to latlong_to_meshcode()")

    # Latitude

    lat_in_min = lat * 60.0

    code12 = int(lat_in_min / 40)  # codeの1, 2文字目部分（数値）
    lat_rest_in_min = lat_in_min - code12 * 40   # 残差

    code5 = int(lat_rest_in_min / 5 )   #  code 5文字目 二次メッシュの１区画は緯度5分
    lat_rest_in_min -= code5 * 5        #  残差

    code7 = int(lat_rest_in_min / (5/10)) # code 7文字目 三次メッシュの１区画は緯度５分の 1/10

    # Longitude

    code34 = int(long) - 100  # codeの3, 4文字目部分（数値）

    long_rest_in_deg = long - int(long)

    code6 = int(long_rest_in_deg * 8)
    long_rest_in_deg -= code6 / 8;

    code8 = int(long_rest_in_deg / (1/80) )

    code = code12 * 100 + code34

    if order >= 2:
        code = code * 100 + code5 * 10 + code6

    if order == 3:
        code = code * 100 + code7 * 10 + code8

    return int(code)

##################
#  度、分、秒を受けとり、度の単位で小数点表示に変換して返す。
#
# Usage:
#  deg = meshlib.dms_to_deg(35, 20, 15)

def dms_to_deg (deg, min, sec):

    in_deg = deg + min / 60.0 + sec / 3600.0

    return in_deg

##################
# 度の単位の小数点表示を受けとり、 度、分、秒に変換して返す。
# 戻り値は数値データが３つ並んだタプル。
#
# Usage:
#  dms = meshlib.deg_to_dms(35.3375)

def deg_to_dms(in_deg):

    deg         = int(in_deg)
    rest_in_deg = in_deg - deg

    min         = int( rest_in_deg * 60 )
    rest_in_min = rest_in_deg * 60 - min;

    sec = rest_in_min * 60;

    # もし繰り上がってしまっていたら、その分の調整、

    if sec >= 60:   # 秒が60を越えたら
        sec -= 60
        min += 1

    if min == 60:    # 分が60 を越えたら
       min -= 60
       deg += 1

    return (deg, min, sec)  # タプルを返す。

def latlong_to_detail_meshcode(lat, log, param=0):
    """
    = (QUOTIENT($A2*60,40) &
        INT($B2-100) &
        [p1] QUOTIENT(MOD($A2*60,40),40/8) &
        [p2] QUOTIENT(MOD(($B2-100),1)*60,60/8) &
        [p3] QUOTIENT(MOD(MOD($A2*60,40),40/8),40/8/10) &
        [p4] QUOTIENT(MOD(MOD(($B2-100),1)*60,60/8),60/8/10) &
        [p5] QUOTIENT(MOD(MOD(MOD($A2*60,40),40/8),40/8/10),40/8/10/2)*2 +
            QUOTIENT(MOD(MOD(MOD(($B2-100),1)*60,60/8),60/8/10),60/8/10/2)+1)*1
    """
    lat60, mod = divmod(lat * 60, 40)
    lng = log - 100

    p1 = divmod(mod, 40/8 )[0]
    p2 = divmod(divmod(lng,1)[1]*60,(60/8))[0]
    p3 = divmod(divmod(mod, 40/8 )[1], 40/8/10)[0]
    p4 = divmod(divmod(divmod(lng,1)[1]*60,(60/8))[1], (60/8/10))[0]
    p5 = (divmod(divmod(divmod(mod, 40/8)[1], 40/8/10)[1], 40/8/10/2 )[0] * 2 + \
          divmod(divmod(divmod(divmod(lng,1)[1]*60,(60/8))[1], (60/8/10))[1],60/8/10/2)[0]) + 1

    print("lat60", int(lat60))
    print("lng", lng)
    print("p1", p1)
    print("p2", p2)
    print("p3", p3)
    print("p4", p4)
    print("p5", p5)
    #print("p5_1", divmod(divmod(divmod(mod, 40/8)[1], 40/8/10)[1], 40/8/10/2 )[0] * 2)
    #print("p5_2", divmod(divmod(divmod(divmod(lng,1)[1]*60,60/8)[1], (60/8/10))[1],60/8/10/2)[0]) + 1)
    result_mesh = []
    if(param == 2 or param == 0):
        result_mesh.append(str(int(lat60))+str(int(lng))+'-'+str(int(p1))+str(int(p2))+str(int(p3))+str(int(p4))+'-'+str(int(p5)))
        if(param == 2 ):
            return result_mesh

    """
        [p6] = QUOTIENT(MOD(MOD(MOD(MOD($A2*60,40),40/8),40/8/10),40/8/10/2),40/8/10/2/2)*2 +
               QUOTIENT(MOD(MOD(MOD(MOD(($B2-100),1)*60,60/8),60/8/10),60/8/10/2),60/8/10/2/2)+1)*1
    """
    p6 = (divmod(divmod(divmod(divmod(mod, 40/8)[1], 40/8/10)[1], 40/8/10/2 )[1],40/8/10/2/2)[0] * 2 + \
          divmod(divmod(divmod(divmod(divmod(lng,1)[1]*60,60/8)[1], 60/8/10)[1],60/8/10/2)[1],60/8/10/2/2)[0] + 1)
    if(param == 4 or param == 0):
        result_mesh.append(str(int(lat60))+str(int(lng))+'-'+str(int(p1))+str(int(p2))+str(int(p3))+str(int(p4))+'-'+str(int(p5))+'-'+str(int(p6)))
        if(param == 4 ):
            return result_mesh
    """
        [p7] = QUOTIENT(MOD(MOD(MOD(MOD(MOD($A2*60,40),40/8),40/8/10),40/8/10/2),40/8/10/2/2),40/8/10/2/2/2)*2 +
               QUOTIENT(MOD(MOD(MOD(MOD(MOD(($B2-100),1)*60,60/8),60/8/10),60/8/10/2),60/8/10/2/2),60/8/10/2/2/2)+1)*1
    """
    p7 = (divmod(divmod(divmod(divmod(divmod(mod, 40/8)[1], 40/8/10)[1], 40/8/10/2 )[1],40/8/10/2/2)[1],40/8/10/2/2/2)[0] * 2 + \
          divmod(divmod(divmod(divmod(divmod(divmod(lng,1)[1]*60,60/8)[1], 60/8/10)[1],60/8/10/2)[1],60/8/10/2/2)[1],60/8/10/2/2/2)[0] + 1)

    result_mesh.append(str(int(lat60))+str(int(lng))+'-'+str(int(p1))+str(int(p2))+str(int(p3))+str(int(p4))+'-'+str(int(p5))+'-'+str(int(p6))+'-'+str(int(p7)))
    return result_mesh


import ase.io
import fire

from ai2_kit.core.log import get_logger
from ai2_kit import res
from ai2_kit.domain.cp2k import dump_coord_n_cell
from typing import Optional, Literal
from ase import Atoms, Atom
from string import Template
import re
import os
import json

logger = get_logger(__name__)

DEEPMD_DEFAULT_TEMPLATE = os.path.join(res.DIR_PATH, 'catalysis/deepmd.json')
MLP_TRAINING_TEMPLATE = os.path.join(res.DIR_PATH, 'catalysis/mlp-training.yml')

CP2K_DEFAULT_TEMPLATE = os.path.join(res.DIR_PATH, 'catalysis/cp2k.inp')

CP2K_SCF_SEMICONDUCTOR = """\
        # CONFIGURATION FOR SEMICONDUCTOR
        &SCF
            SCF_GUESS RESTART
            EPS_SCF 3e-07
            MAX_SCF 50
            &OUTER_SCF
                EPS_SCF 3e-07
                MAX_SCF 20
            &END OUTER_SCF
            &OT
                MINIMIZER CG
                PRECONDITIONER FULL_SINGLE_INVERSE
                ENERGY_GAP 0.1
            &END OT
        &END SCF
        # END CONFIGURATION FOR SEMICONDUCTOR"""

CP2K_SCF_METAL = """\
        # CONFIGURATION FOR METAL
        &SCF
            SCF_GUESS RESTART
            EPS_SCF 3e-07
            MAX_SCF 500
            ADDED_MOS 500
            CHOLESKY INVERSE
            &SMEAR
                METHOD FERMI_DIRAC
                ELECTRONIC_TEMPERATURE [K] 300
            &END SMEAR
            &DIAGONALIZATION
                ALGORITHM STANDARD
            &END DIAGONALIZATION
            &MIXING
                METHOD BROYDEN_MIXING
                ALPHA 0.3
                BETA 1.5
                NBROYDEN 14
            &END MIXING
        &END SCF
        # END CONFIGURATION FOR METAL"""

CP2K_MOTION_TEMPLATE = """\
&MOTION
  &MD
    ENSEMBLE NVT
    STEPS       $steps
    TIMESTEP    $timestep
    TEMPERATURE $temp
    &THERMOSTAT
       TYPE CSVR
       REGION MASSIVE
       &CSVR
          TIMECON [fs] 100.0
       &END
    &END
  &END MD
  &PRINT
   &TRAJECTORY
     &EACH
       MD 1
     &END EACH
   &END TRAJECTORY
   &VELOCITIES
     &EACH
       MD 1
     &END EACH
   &END VELOCITIES
   &FORCES
     &EACH
       MD 1
     &END EACH
   &END FORCES
   &RESTART_HISTORY
     &EACH
       MD 1000
     &END EACH
   &END RESTART_HISTORY
   &RESTART
     BACKUP_COPIES 3
     &EACH
       MD 1
     &END EACH
   &END RESTART
  &END PRINT
&END MOTION"""

CP2K_SCF_TABLE = {
    'metal': CP2K_SCF_METAL,
    'semi': CP2K_SCF_SEMICONDUCTOR,
}

CP2K_ACCURACY_TABLE = {
    'high': {'cutoff': 1000, 'rel_cutoff': 90 },
    'medium': {'cutoff': 800, 'rel_cutoff': 70 },
    'low': {'cutoff': 600, 'rel_cutoff': 50 },
}

class ConfigBuilder:

    def __init__(self):
        self._atoms: Optional[Atoms] = None

    def load_system(self, file: str, **kwargs):
        """
        Loading system to memory using ASE
        """
        assert self._atoms is None, 'atoms is already loaded'
        self._atoms = ase.io.read(file, **kwargs)  # type: ignore
        return self

    def gen_mlp_training_input(self,
                               out_dir: str = 'out',
                               template_file: str = MLP_TRAINING_TEMPLATE):
        # Read yaml file as text so that the comments are preserved
        with open(template_file, 'r') as fp:
            text = fp.read()
        # Generate the type_map and mass_map automatically
        assert self._atoms is not None, 'atoms must be loaded first'
        type_map = list(set(self._atoms.get_chemical_symbols()))
        mass_map = [Atom(symbol).mass for symbol in type_map]

        out_data = Template(text).substitute(
            type_map=json.dumps(type_map),
            mass_map=json.dumps(mass_map),
        )
        os.makedirs(out_dir, exist_ok=True)
        mlp_training_input_path = os.path.join(out_dir, 'mlp-training.yml')
        with open(mlp_training_input_path, 'w', encoding='utf-8') as fp:
            fp.write(out_data)

    def gen_deepmd_input(self,
                         out_dir: str = 'out',
                         steps: int = 10000,
                         template_file: str = DEEPMD_DEFAULT_TEMPLATE):
        with open(template_file, 'r') as fp:
            data = json.load(fp)
        os.makedirs(out_dir, exist_ok=True)

        data['training']['numb_steps'] = steps
        deepmd_input_path = os.path.join(out_dir, 'deepmd.json')
        with open(deepmd_input_path, 'w', encoding='utf-8') as fp:
            json.dump(data, fp, indent=4)

    def gen_plumed_input(self, out_dir: str = 'out'):
        """
        Generate Plumed input files
        Args:
            out_dir: output directory, plumed.dat will be generated in this directory
        """
        assert self._atoms is not None, 'atoms must be loaded first'
        plumed_input = [ 'UNITS LENGTH=A' ]
        # define atoms groups by element, e.g:
        # Ag: GROUP ATOMS=1,2,3,4,5,6,7,8,9,10
        # O: GROUP ATOMS=11,12,13,14,15,16,17,18,19,20
        element2ids = {}
        for i, element in enumerate(self._atoms.get_chemical_symbols(), start=1):
            element2ids.setdefault(element, []).append(i)
        for element, ids in element2ids.items():
            plumed_input.append(f'{element}: GROUP ATOMS={",".join(map(str, ids))}')

        # define reaction coordinate
        plumed_input.extend([
            '#define more groups if needed',
            '',
            '# define reaction coordinates, e.g. CV1, CV2, ...',
            '# you may define as many as you want',
            'CV1:',
            '',
            '# define sampling method: metadynamics',
            'metad: METAD ARG=CV1 SIGMA=0.1 HEIGHT=5 PACE=100 TEMP=1000 FILE=HILLS',
            '# define more commands if you need',
            '',
            '# print CVs',
            'PRINT STRIDE=10 ARG=CV1,metad.bias FILE=COLVAR',
        ])
        os.makedirs(out_dir, exist_ok=True)
        plumed_input_path = os.path.join(out_dir, 'plumed.dat')
        with open(plumed_input_path, 'w', encoding='utf-8') as fp:
            fp.write('\n'.join(plumed_input))

    def gen_lammps_input(self):
        ... # TODO

    def gen_report(self):
        ... # TODO

    def gen_cp2k_input(self,
                       basic_set_file: str = 'BASIS_MOLOPT',
                       potential_file: str = 'GTH_POTENTIALS',
                       out_dir: str = 'out',
                       style: Literal['metal', 'semi'] = 'metal',
                       accuracy: Literal['high', 'medium', 'low'] = 'medium',
                       aimd: bool = False,
                       temp: float = 330.,
                       steps: int = 1000,
                       timestep: float = 0.5,
                       parameter_file: str = 'dftd3.dat',
                       template_file: str = CP2K_DEFAULT_TEMPLATE,):
        """
        Generate CP2K input files

        You should call `load_system` first to load system into memory.
        And you may also need to ensure the basic set and potential files
        you want to use are available in CP2K_DATA_DIR,
        or else you have to specify the full path instead of their file name.

        Args:
            basic_set_file: basic set file, can be path or file in CP2K_DATA_DIR
            potential_file: potential file, can be path or file in CP2K_DATA_DIR
            out_dir: output directory, cp2k.inp and coord_n_cell.inc will be generated in this directory
            style: 'metal' or 'semi'
            accuracy: 'high', 'medium' or 'low'
            aimd: whether to run AIMD
            temp: temperature for AIMD
            steps: steps for AIMD
            timestep: timestep for AIMD
            parameter_file: parameter file, can be path or file in CP2K_DATA_DIR
            template_file: template file, no need to change in most cases
        """

        assert self._atoms is not None, 'atoms must be loaded first'
        # get available basic set and potential
        with open(find_cp2k_data_file(basic_set_file), 'r') as fp:
            basic_set_table = parse_cp2k_basic_set_potential(fp)
        with open(find_cp2k_data_file(potential_file), 'r') as fp:
            potential_table = parse_cp2k_basic_set_potential(fp)
        with open(template_file, 'r') as fp:
            template = fp.read()
        # build kinds
        elements = set(self._atoms.get_chemical_symbols())

        kinds = ''
        total_ve = 0  # Valence Electron

        for element in elements:
            basic_set_all = basic_set_table.get(element)
            basic_set = basic_set_all[-1]
            logger.info(f'BASIC_SET {basic_set} is selected for {element}')

            potential_all = potential_table.get(element)
            potential = next(p for p in potential_all if 'PBE' in p)
            logger.info(f'POTENTIAL {potential} is selected for {element}')
            kinds += '\n'.join([
                f'    &KIND {element}',
                f'        BASIS_SET {basic_set}',
                f'        # All available BASIS_SET:',
                f'        # {" ".join(basic_set_all)}',
                f'        POTENTIAL {potential}',
                f'        # All available POTENTIAL:',
                f'        # {" ".join(potential_all)}',
                f'    &END KIND',
                '', # This empty line is required
            ])
            # get ve from the last number of potential
            # for example, if the potential is GTH-OLYP-q6
            # then the ve is 6
            ve = int(re.match(r'.*?(\d+)$', potential).group(1))
            total_ve += ve * self._atoms.get_chemical_symbols().count(element)
        logger.info("total valence electron of the system: %d" % total_ve)

        motion = ''
        if aimd:
            motion = Template(CP2K_MOTION_TEMPLATE).substitute(
                steps=steps,
                timestep=timestep,
                temp=temp,
            )

        template_vars = {
            'run_type': 'MD' if aimd else 'ENERGY_FORCE',
            'basic_set_file': basic_set_file,
            'potential_file': potential_file,
            'parameter_file': parameter_file,
            'uks': 'T' if total_ve % 2 else 'F',
            'kinds': kinds,
            'motion': motion,
            'scf': CP2K_SCF_TABLE[style],
            **CP2K_ACCURACY_TABLE[accuracy],
        }

        cp2k_input = Template(template).substitute(**template_vars)
        os.makedirs(out_dir, exist_ok=True)
        cp2k_input_path = os.path.join(out_dir, 'cp2k.inp')
        with open(cp2k_input_path, 'w') as fp:
            fp.write(cp2k_input)
        logger.info(f'CP2K input file is generated at {cp2k_input_path}')

        coord_n_cell_path = os.path.join(out_dir, 'coord_n_cell.inc')
        with open(coord_n_cell_path, 'w') as fp:
            dump_coord_n_cell(fp, self._atoms)
        logger.info(f'coord_n_cell.inc is generated at {coord_n_cell_path}')

def find_cp2k_data_file(file: str):
    """
    find CP2K data file in CP2K_DATA_DIR
    """
    if not os.path.exists(file):
        cp2k_data_dir = os.environ.get('CP2K_DATA_DIR')
        if cp2k_data_dir is None:
            raise FileNotFoundError(f'Cannot find {file}')
        file = os.path.join(cp2k_data_dir, file)
        if not os.path.exists(file):
            raise FileNotFoundError(f'Cannot find {file}')
    return file

def parse_cp2k_basic_set_potential(fp):
    """
    get all available basic set and potential from CP2K data file
    """
    parsed = {}
    for line in fp:
        line: str = line.strip()
        if line.startswith('#'):
            continue # skip comment
        tokens = line.split()
        if 0 == len(tokens):
            continue  # skip empty line
        if re.match(r'^[A-Z][a-z]*$', tokens[0]):
            parsed.setdefault(tokens[0], []).append(tokens[1])
    return parsed

class CmdEntries:

    def build_config(self):
        """
        Build config file for catalyst
        """
        return ConfigBuilder

def cli_main():
    fire.Fire(CmdEntries)

if __name__ == '__main__':
    fire.Fire(ConfigBuilder)


import pandas as pd
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score

# Charger les données dans un DataFrame Pandas
df = pd.read_csv('data/LDA_transformed_data.csv')

#Neural network
couches = 80
from sklearn.neural_network import MLPClassifier
X = df.iloc[:,1:].values
y = df.iloc[:,0].values.astype(int)
clf = MLPClassifier(solver='lbfgs', alpha=1e-5, hidden_layer_sizes=(couches,), random_state=42, max_iter=10000)
clf.fit(X, y)
MLPClassifier(alpha=1e-05, hidden_layer_sizes=(couches,), random_state=42, solver='lbfgs')

#Import well_data
df_well = pd.read_csv('data/well_data_with_facies.csv')

#LDA of well_data
# Sélectionner les variables d'entrée et la variable cible
X = df_well[['GR', 'ILD_log10', 'DeltaPHI', 'PHIND', 'PE']]
y = df_well['Facies']

# Transformation des données avec la LDA du training data
data = pd.read_csv('data/training_data.csv')
databis=data.iloc[:,0]

# Sélectionner les variables d'entrée et la variable cible
X_training = data[['GR', 'ILD_log10', 'DeltaPHI', 'PHIND', 'PE']]
y_training = data['Facies']

# Effectuer une AFD
lda = LinearDiscriminantAnalysis()
n=lda.fit(X_training, y_training)

#Transformer les données à prédire avec le même LDA que celui du training
ldafact_well=lda.transform(X)

#Prediction des données
y_pred = clf.predict(ldafact_well)

#print(confusion_matrix(y, y_pred))
print("LDA accuracy: " + str(accuracy_score(y, y_pred)))


"""Binary tree/non-tree classifier(s)."""
import glob
from os import path

import dask
import maxflow as mf
import numpy as np
import rasterio as rio
from dask import diagnostics
from sklearn import ensemble

from . import pixel_features, pixel_response, settings, utils

__all__ = ["ClassifierTrainer", "Classifier"]

MOORE_NEIGHBORHOOD_ARR = np.array([[0, 0, 0], [0, 0, 1], [1, 1, 1]])

class ClassifierTrainer(object):
    """Train binary tree/non-tree classifier(s) of the pixel features."""

    def __init__(
        self,
        *,
        num_estimators=None,
        sigmas=None,
        num_orientations=None,
        neighborhood=None,
        min_neighborhood_range=None,
        num_neighborhoods=None,
        tree_val=None,
        nontree_val=None,
        **adaboost_kws,
    ):
        """
        Initialize the classifier.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for details.

        Parameters
        ----------
        num_estimators : int, optional
            The maximum number of estimators at which boosting is terminated. Directly
            passed to the `n_estimators` keyword argument of
            `sklearn.ensemble.AdaBoostClassifier`. If no value is provided, the default
            value set in `settings.CLF_DEFAULT_NUM_ESTIMATORS` will be taken.
        sigmas : list-like, optional
            The list of scale parameters (sigmas) to build the Gaussian filter bank that
            will be used to compute the pixel-level features. The provided argument will
            be passed to the initialization method of the `PixelFeaturesBuilder` class.
            If no value is provided, the default value set in
            `settings.GAUSS_DEFAULT_SIGMAS` will be taken.
        num_orientations : int, optional
            The number of equally-distributed orientations to build the Gaussian filter
            bank that will be used to compute the pixel-level features. The provided
            argument will be passed to the initialization method of the
            `PixelFeaturesBuilder` class. If no value is provided, the default value set
            in `settings.GAUSS_DEFAULT_NUM_ORIENTATIONS` will be taken.
        neighborhood : array-like, optional
            The base neighborhood structure that will be used to compute the entropy
            features. Theprovided argument will be passed to the initialization method
            of the `PixelFeaturesBuilder` class. If no value is provided, a square with
            a side size of `2 * min_neighborhood_range + 1` will be used.
        min_neighborhood_range : int, optional
            The range (i.e., the square radius) of the smallest neigbhorhood window that
            will be used to compute the entropy features. The provided argument will be
            passed to the initialization method of the `PixelFeaturesBuilder` class. If
            no value is provided, the default value set in
            `settings.ENTROPY_DEFAULT_MIN_NEIGHBORHOOD_RANGE` will be taken.
        num_neighborhoods : int, optional
            The number of neigbhorhood windows (whose size follows a geometric
            progression starting at `min_neighborhood_range`) that will be used to
            compute the entropy features. The provided argument will be passed to the
            initialization method of the `PixelFeaturesBuilder` class. If no value is
            provided, the default value set in
            `settings.ENTROPY_DEFAULT_NUM_NEIGHBORHOODS` will be taken.
        tree_val : int, optional
            The value that designates tree pixels in the response images. The provided
            argument will be passed to the initialization method of the
            `PixelResponseBuilder` class. If no value is provided, the default value set
            in `settings.RESPONSE_DEFAULT_TREE_VAL` will be taken.
        nontree_val : int, optional
            The value that designates non-tree pixels in the response images. The
            provided argument will be passed to the initialization method of the
            `PixelResponseBuilder` class. If no value is provided, the default value set
            in `settings.RESPONSE_DEFAULT_NONTREE_VAL` will be taken.
        adaboost_kws : key-value pairings, optional
            Keyword arguments that will be passed to
            `sklearn.ensemble.AdaBoostClassifier`.
        """
        super(ClassifierTrainer, self).__init__()

        if num_estimators is None:
            num_estimators = settings.CLF_DEFAULT_NUM_ESTIMATORS
        self.num_estimators = num_estimators

        self.pixel_features_builder_kws = dict(
            sigmas=sigmas,
            num_orientations=num_orientations,
            neighborhood=neighborhood,
            min_neighborhood_range=min_neighborhood_range,
            num_neighborhoods=num_neighborhoods,
        )
        self.pixel_response_builder_kws = dict(
            tree_val=tree_val, nontree_val=nontree_val
        )
        self.adaboost_kws = adaboost_kws

    def train_classifier(
        self,
        *,
        split_df=None,
        response_img_dir=None,
        img_filepaths=None,
        response_img_filepaths=None,
        img_dir=None,
        img_filename_pattern=None,
        method=None,
        img_cluster=None,
    ):
        """
        Train a classifier.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for more
        details.

        Parameters
        ----------
        split_df : pandas DataFrame, optional
            Data frame with the train/test split.
        response_img_dir : str representing path to a directory, optional
            Path to the directory where the response tiles are located. Required if
            providing `split_df`. Otherwise `response_img_dir` might either be ignored
            if providing `response_img_filepaths`, or be used as the directory where the
            images whose filename matches `img_filename_pattern` are to be located.
        img_filepaths : list-like, optional
            List of paths to the input tiles whose features will be used to train the
            classifier. Ignored if `split_df` is provided.
        response_img_filepaths : list-like, optional
            List of paths to the binary response tiles that will be used to train the
            classifier. Ignored if `split_df` is provided.
        img_dir : str representing path to a directory, optional
            Path to the directory where the images whose filename matches
            `img_filename_pattern` are to be located. Ignored if `split_df` or
            `img_filepaths` is provided.
        img_filename_pattern : str representing a file-name pattern, optional
            Filename pattern to be matched in order to obtain the list of images. If no
            value is provided, the default value set in
            `settings.IMG_DEFAULT_FILENAME_PATTERN` will be taken. Ignored if `split_df`
            or `img_filepaths` is provided.
        method : {'cluster-I', 'cluster-II'}, optional
            Method used in the train/test split.
        img_cluster : int, optional
            The label of the cluster of tiles. Only used if `method` is 'cluster-II'.

        Returns
        -------
        clf : scikit-learn AdaBoostClassifier
            The trained classifier
        """
        if split_df is None and response_img_filepaths is None:
            # this is the only case that needs argument tweaking: otherwise, if we pass
            # `img_filepaths`/`img_dir` to `build_features` and `response_img_dir` to
            # `build_response`, the latter would build a response with all the image
            # files in `response_img_dir`. Instead, we need to build the response only
            # for the files speficied in `img_filepaths`/`img_dir`
            if img_filepaths is None:
                # TODO: this is copied from `build_features` - ideally, we should DRY it
                if img_filename_pattern is None:
                    img_filename_pattern = settings.IMG_DEFAULT_FILENAME_PATTERN
                if img_dir is None:
                    raise ValueError(
                        "Either `split_df`, `img_filepaths` or `img_dir` must "
                        "be provided"
                    )
                img_filepaths = glob.glob(path.join(img_dir, img_filename_pattern))

            response_img_filepaths = [
                path.join(response_img_dir, path.basename(img_filepath))
                for img_filepath in img_filepaths
            ]

        X = pixel_features.PixelFeaturesBuilder(
            **self.pixel_features_builder_kws
        ).build_features(
            split_df=split_df,
            img_filepaths=img_filepaths,
            img_dir=img_dir,
            img_filename_pattern=img_filename_pattern,
            method=method,
            img_cluster=img_cluster,
        )

        y = pixel_response.PixelResponseBuilder(
            **self.pixel_response_builder_kws
        ).build_response(
            split_df=split_df,
            response_img_dir=response_img_dir,
            response_img_filepaths=response_img_filepaths,
            img_filename_pattern=img_filename_pattern,
            method=method,
            img_cluster=img_cluster,
        )

        clf = ensemble.AdaBoostClassifier(
            n_estimators=self.num_estimators, **self.adaboost_kws
        )
        clf.fit(X, y)

        return clf

    def train_classifiers(self, split_df, response_img_dir):
        """
        Train a classifier for each first-level cluster in `split_df`.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for more
        details.

        Parameters
        ----------
        split_df : pandas DataFrame
            Data frame with the train/test split, which must have an `img_cluster`.
            column with the first-level cluster labels.
        response_img_dir : str representing path to a directory
            Path to the directory where the response tiles are located.

        Returns
        -------
        clf_dict : dictionary
            Dictionary mapping a scikit-learn AdaBoostClassifier to each first-level
            cluster label.
        """
        if "img_cluster" not in split_df:
            raise ValueError(
                "`split_df` must have an 'img_cluster' column ('cluster-II'). "
                "For 'cluster-I', use `train_classifier`."
            )

        clfs_lazy = {}
        for img_cluster, _ in split_df.groupby("img_cluster"):
            clfs_lazy[img_cluster] = dask.delayed(self.train_classifier)(
                split_df=split_df,
                response_img_dir=response_img_dir,
                method="cluster-II",
                img_cluster=img_cluster,
            )

        with diagnostics.ProgressBar():
            clfs_dict = dask.compute(clfs_lazy)[0]

        return clfs_dict

class Classifier(object):
    """Use trained classifier(s) to predict tree pixels."""

    def __init__(
        self,
        *,
        tree_val=None,
        nontree_val=None,
        refine=None,
        refine_beta=None,
        refine_int_rescale=None,
        **pixel_features_builder_kws,
    ):
        """
        Initialize the classifier instance.

        See the `background <https://bit.ly/2KlCICO>`_ example notebook for more
        details.

        Parameters
        ----------
        tree_val : int, optional
            Label used to denote tree pixels in the predicted images. If no value is
            provided, the default value set in `settings.CLF_DEFAULT_TREE_VAL` will be
            taken.
        nontree_val : int, optional
            Label used to denote non-tree pixels in the predicted images. If no value is
            provided, the default value set in `settings.CLF_DEFAULT_NONTREE_VAL` will
            be taken.
        refine : bool, optional
            Whether the pixel-level classification should be refined by optimizing the
            consistence between neighboring pixels. If no value is provided, the default
            value set in `settings.CLF_DEFAULT_REFINE` will be taken.
        refine_beta : int, optional
            Parameter of the refinement procedure that controls the smoothness of the
            labelling. Larger values lead to smoother shapes.  If no value is provided,
            the default value set in `settings.CLF_DEFAULT_REFINE_BETA` will be taken.
        refine_int_rescale : int, optional
            Parameter of the refinement procedure that controls the precision of the
            transformation of float to integer edge weights, required for the employed
            graph cuts algorithm. Larger values lead to greater precision. If no value
            is provided, the default value set in
            `settings.CLF_DEFAULT_REFINE_INT_RESCALE` will be taken.
        pixel_features_builder_kws : dict, optional
            Keyword arguments that will be passed to `detectree.PixelFeaturesBuilder`,
            which customize how the pixel features are built.
        """
        super(Classifier, self).__init__()

        if tree_val is None:
            tree_val = settings.CLF_DEFAULT_TREE_VAL
        if nontree_val is None:
            nontree_val = settings.CLF_DEFAULT_NONTREE_VAL
        if refine is None:
            refine = settings.CLF_DEFAULT_REFINE
        if refine_beta is None:
            refine_beta = settings.CLF_DEFAULT_REFINE_BETA
        if refine_int_rescale is None:
            refine_int_rescale = settings.CLF_DEFAULT_REFINE_INT_RESCALE

        self.tree_val = tree_val
        self.nontree_val = nontree_val
        self.refine = refine
        self.refine_beta = refine_beta
        self.refine_int_rescale = refine_int_rescale

        self.pixel_features_builder_kws = pixel_features_builder_kws

    def classify_img(self, img_filepath, clf, output_filepath=None):
        """
        Use a trained classifier to predict tree pixels in an image.

        Optionally dump the predicted tree/non-tree image to `output_filepath`.

        Parameters
        ----------
        img_filepath : str, file object or pathlib.Path object
            Path to a file, URI, file object opened in binary ('rb') mode, or a Path
            object representing the image to be classified. The value will be passed to
            `rasterio.open`.
        clf : scikit-learn AdaBoostClassifier
            Trained classifier.
        output_filepath : str, file object or pathlib.Path object, optional
            Path to a file, URI, file object opened in binary ('rb') mode, or a Path
            object representing where the predicted image is to be dumped. The value
            will be passed to `rasterio.open` in 'write' mode.

        Returns
        -------
        y_pred : numpy ndarray
            Array with the pixel responses.
        """
        # ACHTUNG: Note that we do not use keyword-only arguments in this method because
        # `output_filepath` works as the only "optional" argument
        src = rio.open(img_filepath)
        img_shape = src.shape

        X = pixel_features.PixelFeaturesBuilder(
            **self.pixel_features_builder_kws
        ).build_features_from_filepath(img_filepath)

        if not self.refine:
            y_pred = clf.predict(X).reshape(img_shape)
        else:
            p_nontree, p_tree = np.hsplit(clf.predict_proba(X), 2)
            g = mf.Graph[int]()
            node_ids = g.add_grid_nodes(img_shape)
            P_nontree = p_nontree.reshape(img_shape)
            P_tree = p_tree.reshape(img_shape)

            # The AdaBoost probabilities are floats between 0 and 1, and the graph cuts
            # algorithm requires an integer representation.  Therefore, we will multiply
            # the probabilities by an arbitrary large number and then transform the
            # result to integers. For instance, we could use a `refine_int_rescale` of
            # `100` so that the probabilities are rescaled into integers between 0 and
            # 100 like percentages). The larger `refine_int_rescale`, the greater the
            # precision.
            # ACHTUNG: the data term when the pixel is a tree is `log(1 - P_tree)`,
            # i.e., `log(P_nontree)`, so the two lines below are correct
            D_tree = (self.refine_int_rescale * np.log(P_nontree)).astype(int)
            D_nontree = (self.refine_int_rescale * np.log(P_tree)).astype(int)
            # TODO: option to choose Moore/Von Neumann neighborhood?
            g.add_grid_edges(
                node_ids, self.refine_beta, structure=MOORE_NEIGHBORHOOD_ARR
            )
            g.add_grid_tedges(node_ids, D_tree, D_nontree)
            g.maxflow()
            # y_pred = g.get_grid_segments(node_ids)
            # transform boolean `g.get_grid_segments(node_ids)` to an array of
            # `self.tree_val` and `self.nontree_val`
            y_pred = np.full(img_shape, self.nontree_val)
            y_pred[g.get_grid_segments(node_ids)] = self.tree_val

        # TODO: make the profile of output rasters more customizable (e.g., via the
        # `settings` module)
        # output_filepath = path.join(output_dir,
        #                             f"tile_{tile_start}-{tile_end}.tif")
        if output_filepath is not None:
            with rio.open(
                output_filepath,
                "w",
                driver="GTiff",
                width=y_pred.shape[1],
                height=y_pred.shape[0],
                count=1,
                dtype=np.uint8,
                nodata=self.nontree_val,
                crs=src.crs,
                transform=src.transform,
            ) as dst:
                dst.write(y_pred.astype(np.uint8), 1)

        src.close()
        return y_pred

    def _classify_imgs(self, img_filepaths, clf, output_dir):
        pred_imgs_lazy = []
        pred_img_filepaths = []
        for img_filepath in img_filepaths:
            # filename, ext = path.splitext(path.basename(img_filepath))
            # pred_img_filepath = path.join(
            #     output_dir, f"{filename}-pred{ext}")
            pred_img_filepath = path.join(output_dir, path.basename(img_filepath))
            pred_imgs_lazy.append(
                dask.delayed(self.classify_img)(img_filepath, clf, pred_img_filepath)
            )
            pred_img_filepaths.append(pred_img_filepath)

        with diagnostics.ProgressBar():
            dask.compute(*pred_imgs_lazy)

        return pred_img_filepaths

    def classify_imgs(
        self,
        split_df,
        output_dir,
        *,
        clf=None,
        clf_dict=None,
        method=None,
        img_cluster=None,
    ):
        """
        Use trained classifier(s) to predict tree pixels in multiple images.

        Use `clf` or `clf_dict` for the classifier(s) depending on the train/ test split
        method, and dump the predicted tree/non-tree images to `output_dir`. See the
        `background <https://bit.ly/2KlCICO>`_ example notebook for more details.

        Parameters
        ----------
        split_df : pandas DataFrame, optional
            Data frame with the train/test split.
        output_dir : str or pathlib.Path object
            Path to the directory where the predicted images are to be dumped.
        clf : scikit-learn AdaBoostClassifier
            Trained classifier.
        clf_dict : dictionary
            Dictionary mapping a trained scikit-learn AdaBoostClassifier to each
            first-level cluster label.
        method : {'cluster-I', 'cluster-II'}, optional
            Method used in the train/test split.
        img_cluster : int, optional
            The label of the cluster of tiles. Only used if `method` is 'cluster-II'.

        Returns
        -------
        pred_imgs : list or dict
            File paths of the dumped tiles.
        """
        if method is None:
            if "img_cluster" in split_df:
                method = "cluster-II"
            else:
                method = "cluster-I"

        if method == "cluster-I":
            if clf is None:
                raise ValueError("If using 'cluster-I' method, `clf` must be provided")
            return self._classify_imgs(
                split_df[~split_df["train"]]["img_filepath"], clf, output_dir
            )
        else:
            if img_cluster is not None:
                if clf is None:
                    if clf_dict is not None:
                        clf = clf_dict[img_cluster]
                    else:
                        raise ValueError("Either `clf` or `clf_dict` must be provided")

                return self._classify_imgs(
                    utils.get_img_filepaths(split_df, img_cluster, False),
                    clf,
                    output_dir,
                )

            if clf_dict is None:
                raise ValueError(
                    "If using 'cluster-II' method and not providing `img_cluster`, "
                    "`clf_dict` must be provided"
                )
            pred_imgs = {}
            for img_cluster, _ in split_df.groupby("img_cluster"):
                pred_imgs[img_cluster] = self._classify_imgs(
                    utils.get_img_filepaths(split_df, img_cluster, False),
                    clf_dict[img_cluster],
                    output_dir,
                )

            return pred_imgs