# cif_generator.py
"""
テンプレートと設定情報に基づくCIFファイル生成モジュール

[概要]
このスクリプトは、マスタープログラム (`run_workflow.py`) から呼び出されることを
前提とした「部品（モジュール）」です。
`run_workflow.py` から渡された設定情報（テンプレートのパス、置換する元素、
価数、組成比など）に基づいて、最終的なCIFファイルを一括で生成します。

[主な機能]
- テンプレート (.j2) とメタデータ (.json) を使用したCIFファイルの生成。
- 置換後の化合物が電荷的に中性かどうかの簡易的な検証。
- 生成日時や実験名に基づいた、整理された出力フォルダの作成。

[使い方]
このスクリプトは通常、直接実行しません。`run_workflow.py`が内部で
`CifGenerator`クラスをインポートして使用します。
"""
import os
import itertools
from datetime import datetime
from jinja2 import Environment, FileSystemLoader

class CifGenerator:
    # Class content is unchanged...
    ELEMENT_DATA = {
        'H': 'Hydrogen', 'He': 'Helium', 'Li': 'Lithium', 'Be': 'Beryllium', 'B': 'Boron',
        'C': 'Carbon', 'N': 'Nitrogen', 'O': 'Oxide', 'F': 'Fluoride', 'Ne': 'Neon',
        'Na': 'Sodium', 'Mg': 'Magnesium', 'Al': 'Aluminum', 'Si': 'Silicon', 'P': 'Phosphide',
        'S': 'Sulfide', 'Cl': 'Chloride', 'Ar': 'Argon', 'K': 'Potassium', 'Ca': 'Calcium',
        'Sc': 'Scandium', 'Ti': 'Titanium', 'V': 'Vanadium', 'Cr': 'Chromium', 'Mn': 'Manganese',
        'Fe': 'Iron', 'Co': 'Cobalt', 'Ni': 'Nickel', 'Cu': 'Copper', 'Zn': 'Zinc',
        'Ga': 'Gallium', 'Ge': 'Germanium', 'As': 'Arsenide', 'Se': 'Selenide', 'Br': 'Bromide',
        'Kr': 'Krypton', 'Rb': 'Rubidium', 'Sr': 'Strontium', 'Y': 'Yttrium', 'Zr': 'Zirconium',
        'Nb': 'Niobium', 'Mo': 'Molybdenum', 'Tc': 'Technetium', 'Ru': 'Ruthenium',
        'Rh': 'Rhodium', 'Pd': 'Palladium', 'Ag': 'Silver', 'Cd': 'Cadmium', 'In': 'Indium',
        'Sn': 'Tin', 'Sb': 'Antimonide', 'Te': 'Telluride', 'I': 'Iodide', 'Xe': 'Xenon',
        'Cs': 'Caesium', 'Ba': 'Barium', 'La': 'Lanthanum', 'Ce': 'Cerium', 'Pr': 'Praseodymium',
        'Nd': 'Neodymium', 'Pm': 'Promethium', 'Sm': 'Samarium', 'Eu': 'Europium',
        'Gd': 'Gadolinium', 'Tb': 'Terbium', 'Dy': 'Dysprosium', 'Ho': 'Holmium',
        'Er': 'Erbium', 'Tm': 'Thulium', 'Yb': 'Ytterbium', 'Lu': 'Lutetium',
        'Hf': 'Hafnium', 'Ta': 'Tantalum', 'W': 'Tungsten', 'Re': 'Rhenium',
        'Os': 'Osmium', 'Ir': 'Iridium', 'Pt': 'Platinum', 'Au': 'Gold', 'Hg': 'Mercury',
        'Tl': 'Thallium', 'Pb': 'Lead', 'Bi': 'Bismuthide', 'Po': 'Polonium', 'At': 'Astatine',
        'Rn': 'Radon', 'Fr': 'Francium', 'Ra': 'Radium'
    }
    
    def __init__(self, config):
        self.config = config
        self.template_file = config['template_file']
        self.jinja_env = Environment(loader=FileSystemLoader(os.path.dirname(self.template_file) or '.'), trim_blocks=True)

    def run(self):
        """Runs the CIF generation process based on the config."""
        template = self.jinja_env.get_template(os.path.basename(self.template_file))
        
        output_dir = self._create_output_dir()
        
        substitutions = self.config['substitutions']
        roles = self.config['roles']
        stoichiometry = self.config['stoichiometry']
        
        role_names = list(substitutions.keys())
        element_lists = [self._process_sub_list(substitutions[r]) for r in role_names]
        all_combinations = list(itertools.product(*element_lists))

        print(f"ℹ️  {len(all_combinations)}個の置換CIFファイルを生成します...")
        
        for i, combination in enumerate(all_combinations):
            context = {}
            formula_parts = {}
            
            for j, role_name in enumerate(role_names):
                element = combination[j]
                context[role_name] = {
                    'symbol': element['symbol'],
                    'valence': element.get('valence', roles[role_name]['valence']),
                    'name': self.ELEMENT_DATA.get(element['symbol'], element['symbol'])
                }
                formula_parts[role_name] = {'symbol': element['symbol'], 'count': stoichiometry[role_name]}

            self._validate_charge(context, stoichiometry)
            self._populate_formulas(context, formula_parts)
            
            final_cif_content = self._create_cif_content(template, context)
            
            cation_keys = sorted([k for k in formula_parts if k.startswith('cation')])
            anion_keys = sorted([k for k in formula_parts if k.startswith('anion')])
            ordered_keys = cation_keys + anion_keys

            filename_parts = []
            for key in ordered_keys:
                part = formula_parts[key]
                count = part['count'] if part['count'] > 1 else ''
                filename_parts.append(f"{part['symbol']}{count}")
            
            filename = "".join(filename_parts) + ".cif"
            filepath = os.path.join(output_dir, filename)

            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(final_cif_content)
        
        print(f"✅ 成功: 合計 {len(all_combinations)} 個のCIFファイルを '{output_dir}' に生成しました。")

    def _process_sub_list(self, sub_list):
        """Processes user-defined substitution list."""
        processed = []
        for item in sub_list:
            if isinstance(item, (list, tuple)):
                processed.append({'symbol': item[0], 'valence': item[1]})
            else:
                processed.append({'symbol': item})
        return processed

    def _validate_charge(self, context, stoichiometry):
        """Calculates total charge and prints a warning if not neutral."""
        total_charge = 0
        for role_name, stoich_count in stoichiometry.items():
            valence = context[role_name]['valence']
            total_charge += valence * stoich_count
        
        if total_charge != 0:
            compound = "".join(f"{context[rn]['symbol']}{stoichiometry[rn]}" for rn in stoichiometry)
            print(f"⚠️  警告: 化合物 {compound} の電荷が中性ではありません (総電荷: {total_charge})。")

    def _populate_formulas(self, context, formula_parts):
        """Generates formula strings and adds them to the context."""
        sum_parts = sorted(formula_parts.values(), key=lambda x: x['symbol'])
        context['chemical_formula_sum'] = ' '.join(f"{p['symbol']}{p['count']}" for p in sum_parts)
        
        cation_keys = sorted([k for k in formula_parts if k.startswith('cation')])
        anion_keys = sorted([k for k in formula_parts if k.startswith('anion')])
        ordered_keys = cation_keys + anion_keys
        
        structural_parts = []
        for key in ordered_keys:
            part = formula_parts[key]
            count = part['count'] if part['count'] > 1 else ''
            structural_parts.append(f"{part['symbol']}{count}")
        context['chemical_formula_structural'] = ' '.join(structural_parts)

        context['chemical_name_common'] = ' '.join(self.ELEMENT_DATA.get(formula_parts[key]['symbol'], formula_parts[key]['symbol']) for key in ordered_keys)
    
    def _create_cif_content(self, template, context):
        """Creates the full CIF content including the header."""
        header_lines = [
            "#",
            f"# CIF file generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
            f"# Source Template: {self.template_file}",
            "# Substitutions:",
        ]
        for role_name, role_data in context.items():
            if isinstance(role_data, dict) and 'symbol' in role_data:
                 header_lines.append(f"#   - Role '{role_name}' -> '{role_data['symbol']}' (Valence: {role_data['valence']})")
        header_lines.append("#")
        generation_header = "\n".join(header_lines) + "\n\n"
        
        rendered_cif = template.render(context)
        return generation_header + rendered_cif

    def _create_output_dir(self):
        """Creates a unique output directory for the experiment."""
        base_name = self.config.get('experiment_name', 'cif_generation')
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        output_dir = f"{base_name}_{timestamp}"
        os.makedirs(output_dir, exist_ok=True)
        print(f"📂 出力先フォルダ: '{output_dir}'")
        return output_dir