import json
import requests
import zipfile
import os
import periodictable
import numpy as np



def process_chemical_formula(formula):
    """
    Process a chemical formula to generate a list with the following:
    - Original chemical formula.
    - Chemical formula without numbers.
    - Chemical formula with elements in alphabetical order.
    - Chemical formula in alphabetical order without numbers.
    """
    import re

    # Original formula
    original_formula = formula

    # Formula without numbers
    formula_without_numbers = re.sub(r'\d+', '', formula)

    # Extract elements and their counts
    elements = re.findall(r'([A-Z][a-z]*)(\d*)', formula)
    elements_dict = {element: int(count) if count else 1 for element, count in elements}

    # Formula with elements in alphabetical order
    sorted_elements = sorted(elements_dict.items())
    alphabetical_formula = ''.join([f"{el}{count if count > 1 else ''}" for el, count in sorted_elements])

    # Alphabetical formula without numbers
    alphabetical_formula_without_numbers = ''.join([el for el, _ in sorted_elements])

    return [original_formula, formula_without_numbers, alphabetical_formula, alphabetical_formula_without_numbers]



def read_json(file_name):
    """
    just o read json in one line
    :param file_name:
    :return:
    """
    with open(file_name, 'r') as file:
        data = json.load(file)
    return data

def ndarray_decoder(obj):
    if "__ndarray__" in obj:
        data_type = obj["__ndarray__"][1]
        data = obj["__ndarray__"][2]
        shape = obj["__ndarray__"][0]

        # Handle complex data types
        if data_type == 'complex128':
            # Pair consecutive elements as real and imaginary parts
            complex_data = [complex(data[i], data[i + 1]) for i in range(0, len(data), 2)]
            return np.array(complex_data, dtype='complex128').reshape(shape)
        else:
            return np.array(data, dtype=data_type).reshape(shape)
    return obj


def request_zip(url, local_path,name="downloaded_archive.zip"):
    response = requests.get(url)
    print(response)
    if response.status_code == 200:
        zip_path = os.path.join(local_path, name)
        with open(zip_path, "wb") as file:
            file.write(response.content)

        # Assuming it's a ZIP file, extract it
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(local_path)

        print("Folder/Archive downloaded and extracted.")
    else:
        print("Failed to download the file. Status code:", response.status_code)


def structure_ato_list(data):
    atomic_numbers = data['atoms']
    positions = data['positions']
    result = []

    for atomic_number, position in zip(atomic_numbers, positions):
        # Get the element object from the periodictable module
        element = periodictable.elements[atomic_number]
        # Append the symbol and position to the result list
        result.append([element.symbol, np.array(position)])

    return result


def create_folder(folder_path):
    if not os.path.exists(folder_path):
        # Create the directory
        os.makedirs(folder_path)
        print(f"Folder '{folder_path}' created.")
    else:
        print(f"Folder '{folder_path}' already exists.")


def download_file(url, directory, filename):
    # Send a GET request to the URL
    response = requests.get(url)

    # Check if the request was successful
    if response.status_code == 200:
        # Create the full path for the file
        filepath = f"{directory}/{filename}"

        # Write the content to a file
        with open(filepath, 'wb') as file:
            file.write(response.content)
        print(f"File downloaded successfully: {filepath}")
    else:
        print(f"Failed to download the file. Status code: {response.status_code}")


def convert_numpy(obj):
    """
    Converts objects of type numpy.ndarray to lists for JSON serialization.
    """
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    raise TypeError(f"Object of type '{type(obj).__name__}' is not JSON serializable")

def save_dict_as_json(data, filename):
    try:
        with open(filename, 'w') as file:
            # Use the 'default' argument to specify a function that converts numpy arrays to lists
            json.dump(data, file, indent=4, default=convert_numpy)
        print(f"Dictionary saved as JSON in {filename}")
    except Exception as e:
        print(f"Error saving dictionary as JSON: {e}")


def list_directory_contents(directory):
    try:
        # List all files and directories in the specified path
        contents = os.listdir(directory)
        return contents
    except FileNotFoundError:
        return f"Directory not found: {directory}"
    except PermissionError:
        return f"Permission denied to access: {directory}"


def structure_to_xyz(structure, file_name):
    # Convert cell to lattice string
    cell = np.array(structure["cell"])
    a, b, c = cell[0], cell[1], cell[2]
    lattice = f'Lattice="{a[0]} {a[1]} {a[2]} {b[0]} {b[1]} {b[2]} {c[0]} {c[1]} {c[2]}"'

    # Start building the XYZ file content
    xyz_content = f"{len(structure['atoms'])}\n"
    xyz_content += f"{lattice}  pbc=\"T T F\"\n"

    # Add atoms and their positions
    for atom_num, pos in zip(structure["atoms"], structure["positions"]):
        element = periodictable.elements[atom_num]
        symbol = element.symbol
        xyz_content += f"{symbol} {' '.join(map(str, pos))} 0.00000000 0.00000000 0.00000000\n"  # Add placeholders for magnetic moments and forces

    # Save to a file
    with open(file_name, "w") as file:
        file.write(xyz_content)
    print("Content:")
    print(xyz_content)
    print("XYZ file created: output.xyz")