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 b KQkq e3 0 1",
    "rnbqkbnr/pp1ppppp/8/2p5/4P3/8/PPPP1PPP/RNBQKBNR w KQkq c6 0 2",
    "rnbqkbnr/pp1ppppp/8/2p5/4P3/5N2/PPPP1PPP/RNBQKB1R b KQkq - 1 2",
    "r1bqkbnr/pp1ppppp/2n5/2p5/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3",
    "r1bqkbnr/pp1ppppp/2n5/2p5/2B1P3/5N2/PPPP1PPP/RNBQK2R b KQkq - 3 3",
    "r1bqkbnr/pp1p1ppp/2n5/4p3/2B1P3/5N2/PPPP1PPP/RNBQK2R w KQkq e6 0 4",
    "r1bqkbnr/pp1p1ppp/2n5/4P3/2B5/5N2/PPPP1PPP/RNBQK2R b KQkq - 0 4",
    "rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1",
    "rnbqkb1r/pppppppp/7n/8/4P3/8/PPPP1PPP/RNBQKBNR w KQkq - 1 2",
    "rnbqkb1r/pppppppp/7n/N7/P3P3/P7/PBPPBPPP/RNQQKBRR b KQkq - 0 10"
]

# Создаем массив векторов из ваших данных
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, ax = plt.subplots()

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

# Отображаем выбранные векторы на графике
for vector in selected_vectors:
    ax.scatter(vector[0], vector[1], 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()