"""
(Unofficial) SAD to XSuite Converter: Output Writer - Octupoles
=============================================
Author(s):  John P T Salvesen
Email:      john.salvesen@cern.ch
Date:       09-10-2025
"""

################################################################################
# Import Packages
################################################################################
import xtrack as xt
import xdeps as xd
import numpy as np

from ._000_helpers import extract_multipole_information, \
    generate_magnet_for_replication_names, check_is_simple_quad_sext_oct, \
    check_is_skew_quad_sext_oct
from ..types import ConfigLike

################################################################################
# Lattice File
################################################################################
def create_octupole_lattice_file_information(
        line:       xt.Line,
        line_table: xd.table.Table,
        config:     ConfigLike) -> str:
    """
    Docstring for create_octupole_lattice_file_information
    
    :param line: Description
    :type line: xt.Line
    :param line_table: Description
    :type line_table: xd.table.Table
    :param config: Description
    :type config: ConfigLike
    :return: Description
    :rtype: str
    """

    ########################################
    # Get information
    ########################################
    octs, unique_oct_names = extract_multipole_information(
        line        = line,
        line_table  = line_table,
        mode        = "Octupole")

    oct_lengths    = np.array(sorted(octs.keys()))
    oct_names      = generate_magnet_for_replication_names(octs, "oct")

    ########################################
    # Ensure there are octupoles in the line
    ########################################
    if len(unique_oct_names) == 0:
        return ""

    ########################################
    # Create Output string
    ########################################
    output_string   = """
############################################################
# Octupoles
############################################################
"""

    ########################################
    # Create base elements
    ########################################
    output_string += """
########################################
# Base Elements
########################################"""

    for oct_name, oct_length in zip(oct_names, oct_lengths):
        output_string += f"""
env.new(name = '{oct_name}', parent = xt.Octupole, length = {oct_length})"""

    output_string += "\n"

    ########################################
    # Clone Elements
    ########################################
    output_string += """
########################################
# Cloned Elements
########################################"""

    for oct, oct_length in zip(oct_names, oct_lengths):
        for replica_name in octs[oct_length]:

            # Remove the minus sign if no non minus version exists
            if replica_name.startswith("-"):
                root_name   = replica_name[1:]
                if root_name not in octs[oct_length]:
                    replica_name        = root_name

            if check_is_simple_quad_sext_oct(line, replica_name, "Octupole"):

                if not check_is_skew_quad_sext_oct(line, replica_name, "Octupole"):
                    output_string += f"""
env.new(name = '{replica_name}', parent = '{oct}', k3 = 'k3_{replica_name}')"""
                else:
                    output_string += f"""
env.new(name = '{replica_name}', parent = '{oct}', k3s = 'k3s_{replica_name}')"""

            else:
                # Get the replica information
                k3          = line[replica_name].k3
                k3s         = line[replica_name].k3s
                shift_x     = line[replica_name].shift_x
                shift_y     = line[replica_name].shift_y
                rot_s_rad   = line[replica_name].rot_s_rad

                # Basic information
                oct_generation = f"""
env.new(
    name        = '{replica_name}',
    parent      = '{oct}'"""

                # Strength information
                if k3 != 0:
                    oct_generation += f""",
    k3          = 'k3_{replica_name}'"""
                if k3s != 0:
                    oct_generation += f""",
    k3s         = 'k3s_{replica_name}'"""

                # Misalignments
                if shift_x != 0:
                    oct_generation += f""",
    shift_x     = '{shift_x}'"""
                if shift_y != 0:
                    oct_generation += f""",
    shift_y     = '{shift_y}'"""
                if rot_s_rad != 0:
                    oct_generation += f""",
    rot_s_rad   = '{rot_s_rad}'"""

                # Close the element definition
                oct_generation += """)"""

                # Write to the file
                output_string += oct_generation

    ########################################
    # Return
    ########################################
    output_string += "\n"
    return output_string

################################################################################
# Optics File
################################################################################
def create_octupole_optics_file_information(
        line:       xt.Line,
        line_table: xd.table.Table,
        config:     ConfigLike) -> str:
    """
    Docstring for create_octupole_optics_file_information
    
    :param line: Description
    :type line: xt.Line
    :param line_table: Description
    :type line_table: xd.table.Table
    :param config: Description
    :type config: ConfigLike
    :return: Description
    :rtype: str
    """

    ########################################
    # Get information
    ########################################
    _, unique_oct_names = extract_multipole_information(
        line        = line,
        line_table  = line_table,
        mode        = "Octupole")

    ########################################
    # Ensure there are octupoles in the line
    ########################################
    if len(unique_oct_names) == 0:
        return ""

    ########################################
    # Create Output string
    ########################################
    output_string = """
    ############################################################
    # Octupoles
    ############################################################"""

    for oct in unique_oct_names:
        k3          = None
        k3s         = None

        try:
            k3  = line[oct].k3
        except KeyError:
            try:
                k3  = line[f"-{oct}"].k3
            except KeyError:
                raise KeyError(f"Could not find oct variable {oct} or -{oct} in line.")

        try:
            k3s = line[oct].k3s
        except KeyError:
            try:
                k3s = line[f"-{oct}"].k3s
            except KeyError:
                raise KeyError(f"Could not find oct variable {oct} or -{oct} in line.")

        if k3 == 0:
            k3 = None
        if k3s == 0:
            k3s = None

        if k3 is not None:
            output_string += f"""
    {f'k3_{oct}'}{' ' * (config.OUTPUT_STRING_SEP - len(f'k3_{oct}') + 4)}{'= '}{k3:.24f},"""
        if k3s is not None:
            output_string += f"""
    {f'k3s_{oct}'}{' ' * (config.OUTPUT_STRING_SEP - len(f'k3s_{oct}') + 4)}{'= '}{k3s:.24f},"""

    ########################################
    # Return
    ########################################
    output_string += "\n"
    return output_string
