import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import networkx as nx
from scipy.spatial import distance
from sklearn.decomposition import PCA

# Функция для преобразования шахматной нотации в вектор
def notation_to_vector(notation):
    piece_dict = {'P': 1, 'N': 2, 'B': 3, 'R': 4, 'Q': 5, 'K': 6}
    board_vector = [0]*64
    rows = notation.split('/')
    for i, row in enumerate(rows):
        col = 0
        for char in row:
            if char.isdigit():
                col += int(char)
            else:
                if char.isupper():
                    board_vector[i*8 + col] = piece_dict[char]
                else:
                    board_vector[i*8 + col] = -piece_dict[char.upper()]
                col += 1
    return np.array(board_vector)

# Ваши данные
data = [
    "rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR",
    "rnbqkbnr/pp1ppppp/8/2p5/4P3/8/PPPP1PPP/RNBQKBNR",
    "rnbqkbnr/pp1ppppp/8/2p5/4P3/5N2/PPPP1PPP/RNBQKB1R",
    "r1bqkbnr/pp1ppppp/2n5/2p5/4P3/5N2/PPPP1PPP/RNBQKB1R",
    "r1bqkbnr/pp1ppppp/2n5/2p5/2B1P3/5N2/PPPP1PPP/RNBQK2R",
    "r1bqkbnr/pp1p1ppp/2n1p3/2p5/2B1P3/5N2/PPPP1PPP/RNBQK2R",
    "r1bqkbnr/pp1p1ppp/2n1p3/2p5/2B1P3/2P2N2/PP1P1PPP/RNBQK2R",
    "r1bqkbnr/p2p1ppp/2n1p3/1pp5/2B1P3/2P2N2/PP1P1PPP/RNBQK2R",
    "r1bqkbnr/p2p1ppp/2n1p3/1pp5/4P3/1BP2N2/PP1P1PPP/RNBQK2R",
    "r1bqkbnr/p2p1ppp/2n1p3/1p6/2p1P3/1BP2N2/PP1P1PPP/RNBQK2R",
    "r1bqkbnr/p2p1ppp/2n1p3/1p6/2p1P3/2P2N2/PPBP1PPP/RNBQK2R",
    "r1bqkbnr/3p1ppp/2n1p3/pp6/2p1P3/2P2N2/PPBP1PPP/RNBQK2R",
    "r1bqkbnr/3p1ppp/2n1p3/pp6/2pPP3/2P2N2/PPB2PPP/RNBQK2R",
    "r1bqkbnr/3p1ppp/2n1p3/pp6/4P3/2Pp1N2/PPB2PPP/RNBQK2R",
    "r1bqkbnr/3p1ppp/2n1p3/pp6/4P3/2PQ1N2/PPB2PPP/RNB1K2R",
    "r1bqkb1r/3p1ppp/2n1pn2/pp6/4P3/2PQ1N2/PPB2PPP/RNB1K2R",
    "r1bqkb1r/3p1ppp/2n1pn2/pp2P3/8/2PQ1N2/PPB2PPP/RNB1K2R",
    "r1bqkb1r/3p1ppp/2n1p3/pp1nP3/8/2PQ1N2/PPB2PPP/RNB1K2R",
    "r1bqkb1r/3p1ppp/2n1p3/pp1nP1B1/8/2PQ1N2/PPB2PPP/RN2K2R",
    "r1b1kb1r/2qp1ppp/2n1p3/pp1nP1B1/8/2PQ1N2/PPB2PPP/RN2K2R",
    "r1b1kb1r/2qp1ppp/2n1p3/pp1nP1B1/8/2PQ1N2/PPBN1PPP/R3K2R",
    "r1b1kb1r/2qp1pp1/2n1p2p/pp1nP1B1/8/2PQ1N2/PPBN1PPP/R3K2R",
    "r1b1kb1r/2qp1pp1/2n1p2p/pp1nP3/7B/2PQ1N2/PPBN1PPP/R3K2R",
    "r3kb1r/2qp1pp1/b1n1p2p/pp1nP3/7B/2PQ1N2/PPBN1PPP/R3K2R",
    "r3kb1r/2qp1pp1/b1n1p2p/pp1nP3/7B/1PPQ1N2/P1BN1PPP/R3K2R",
    "r3kb1r/2qp1pp1/b1n1p2p/pp2P3/5n1B/1PPQ1N2/P1BN1PPP/R3K2R"
]

# Создаем массив векторов из ваших данных
vectors = [notation_to_vector(notation.split()[0]) for notation in data]

# Выводим все векторы
for i, vector in enumerate(vectors):
    print(f'Вектор {i+1}:')
    print(vector)
    print()

# Выбираем три случайные точки (вектора)
indices = np.random.choice(len(vectors), 3, replace=False)
selected_vectors = [vectors[i] for i in indices]

# Создаем граф
G = nx.Graph()
for i in range(3):
    G.add_node(i)

# Добавляем ребра так, чтобы получилось дерево (без циклов)
G.add_edge(0, 1)
G.add_edge(0, 2)

# Расчет энергии вершины
def calculate_vertex_energy(vertex, vectors, k):
    energy = 0
    for vector in vectors:
        energy += np.linalg.norm(vertex - vector)**2
    return k * energy

# Расчет энергии ребра
def calculate_edge_energy(edge, vectors, k):
    v1, v2 = vectors[edge[0]], vectors[edge[1]]
    return k * np.linalg.norm(v1 - v2)**2

# Расчет энергии клина
def calculate_wedge_energy(wedge, vectors, k):
    v1, v2, v3 = vectors[wedge[0]], vectors[wedge[1]], vectors[wedge[2]]
    return k * np.linalg.norm(v1 + v3 - 2*v2)**2

# Расчет энергии графа
def calculate_energy(graph, vectors):
    vertex_k = 0.5 # Замените на ваше значение для коэффициента аппроксимации вершин
    edge_k = 0.5 # Замените на ваше значение для коэффициента растяжения ребер
    wedge_k = 0.5 # Замените на ваше значение для коэффициента изгиба клиньев

    energy = 0
    for node in graph.nodes:
        energy += calculate_vertex_energy(vectors[node], vectors, vertex_k)
    for edge in graph.edges:
        energy += calculate_edge_energy(edge, vectors, edge_k)

    # Добавляем энергию клина (только если в графе есть три вершины)
    if len(graph.nodes) == 3:
        energy += calculate_wedge_energy(graph.nodes, vectors, wedge_k)

    return energy

#energy = calculate_energy(G, selected_vectors)
#print(f'Энергия графа: {energy}')

# Применяем PCA
#pca = PCA(n_components=2)
#projected_vectors = pca.fit_transform(vectors)

# Визуализация проекции графа и векторов
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Отображаем все векторы на графике
for vector in vectors:
    ax.scatter(vector[0], vector[1], vector[2], color='b')

# Отображаем выбранные векторы на графике
for vector in selected_vectors:
    ax.scatter(vector[0], vector[1], vector[2], color='r')

# Добавляем ребра между вершинами
for edge in G.edges:
    ax.plot([selected_vectors[edge[0]][0], selected_vectors[edge[1]][0]],
            [selected_vectors[edge[0]][1], selected_vectors[edge[1]][1]],
            [selected_vectors[edge[0]][2], selected_vectors[edge[1]][2]], 'r-')

plt.show()