import os
import math
import numpy as np
import matplotlib.pyplot as plt
from math import sin, pi, sqrt
from collections import defaultdict

# ===============================
# CONFIGURAÇÕES FUNDAMENTAIS
# ===============================

PHI = (1 + sqrt(5)) / 2
MOD = 42
BRACOS_42 = list(range(42))
# ===============================
# UTILIDADES
# ===============================

def is_prime(n):
    if n < 2:
        return False
    for i in range(2, int(sqrt(n)) + 1):
        if n % i == 0:
            return False
    return True

def helicoidal_fn(n, r):
    x = (n - r) / MOD
    return sin(2 * pi * PHI * x) ** 2

# ===============================
# ENTRADA DO USUÁRIO
# ===============================

N = int(input("Limite máximo N (ex: 100000): "))
IGNORE_PRIMES = input("Ignorar primos? (s/n): ").lower() == "s"
EXPORT_TXT = input("Exportar dados TXT? (s/n): ").lower() == "s"

# ===============================
# PREPARAÇÃO DE PASTAS
# ===============================

BASE_DIR = "SACRED_TOTIENT_V2_OUTPUT"
SUBDIRS = ["plots_2D", "plots_3D", "plots_fractal", "data_txt"]

for d in SUBDIRS:
    os.makedirs(os.path.join(BASE_DIR, d), exist_ok=True)

# ===============================
# CÁLCULO PRINCIPAL
# ===============================

dados = []

por_braco = defaultdict(list)

for n in range(1, N + 1):
    r = n % 42            # braço geométrico real (0..41)
    fn = helicoidal_fn(n, r)

    theta = 2 * pi * r / 42
    raio = fn
    z = n

    dados.append((n, r, fn, theta, raio, z))
    por_braco[r].append((n, fn, theta, raio, z))


# ===============================
# GRÁFICOS 2D GERAIS
# ===============================

ns = [d[0] for d in dados]
fns = [d[2] for d in dados]

plt.figure()
plt.scatter(ns, fns, s=1)
plt.title("Distribuição Global Fn")
plt.xlabel("n")
plt.ylabel("Fn")
plt.savefig(f"{BASE_DIR}/plots_2D/Fn_global.png")
plt.close()

# ===============================
# GRÁFICOS POR BRAÇO
# ===============================

for r, vals in por_braco.items():
    ns_r = [v[0] for v in vals]
    fns_r = [v[1] for v in vals]

    plt.figure()
    plt.scatter(ns_r, fns_r, s=2)
    plt.title(f"Fn - Braço r={r}")
    plt.xlabel("n")
    plt.ylabel("Fn")
    plt.savefig(f"{BASE_DIR}/plots_2D/Fn_braco_{r}.png")
    plt.close()

# ===============================
# HÉLICE 3D GLOBAL
# ===============================

from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

xs = [d[3] * math.cos(d[3]) for d in dados]
ys = [d[3] * math.sin(d[3]) for d in dados]
zs = [d[5] for d in dados]

ax.scatter(xs, ys, zs, s=1)
ax.set_title("Hélice Helicoidal Global Fn")
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("n")

plt.savefig(f"{BASE_DIR}/plots_3D/helice_global.png")
plt.close()

# ===============================
# FRACTAL – ZOOM LOCAL
# ===============================

zoom_lim = min(N, 230539333248)
zoom_ns = ns[:zoom_lim]
zoom_fns = fns[:zoom_lim]

plt.figure()
plt.scatter(zoom_ns, zoom_fns, s=2)
plt.title("Zoom Fractal Inicial Fn")
plt.xlabel("n")
plt.ylabel("Fn")
plt.savefig(f"{BASE_DIR}/plots_fractal/zoom_fractal_inicial.png")
plt.close()

# ===============================
# EXPORTAÇÃO TXT
# ===============================

if EXPORT_TXT:
    with open(f"{BASE_DIR}/data_txt/dados_completos.txt", "w") as f:
        f.write("n,r,Fn,theta,raio,z\n")
        for d in dados:
            f.write(",".join(map(str, d)) + "\n")

print("\n=== SACRED TOTIENT FN HELICOIDAL V2 FINALIZADO ===")
print(f"Total de pontos analisados: {len(dados)}")
print("Resultados salvos em:", BASE_DIR)