"""
Агенты-дроны с нейросетями, долгоживучестью и возможностью переквалификации.
"""
import torch
import torch.nn as nn
import numpy as np

class DroneAgent:
    """
    Один дрон как обучаемый узел с персональной моделью.
    """
    def __init__(self, agent_id, state_dim=10, action_dim=4):
        self.id = agent_id
        self.lifetime = 0  # число выживших эпизодов
        self.experience_buffer = []  # (state, action, reward, next_state)
        self.brain = DronePolicyNet(state_dim, action_dim)
        self.optimizer = torch.optim.Adam(self.brain.parameters(), lr=0.001)

    def act(self, state_vec):
        state_t = torch.FloatTensor(state_vec).unsqueeze(0)
        probs = self.brain(state_t)
        action = torch.multinomial(probs, 1).item()
        return action, probs

    def store_experience(self, s, a, r, ns):
        self.experience_buffer.append((s, a, r, ns))
        if len(self.experience_buffer) > 1000:
            self.experience_buffer.pop(0)

    def update_policy(self, gamma=0.99):
        """Простейший policy gradient с учётом lifetime reward."""
        if len(self.experience_buffer) < 32:
            return
        # Здесь должен быть REINFORCE или A2C – в training.py полноценно
        pass

class DronePolicyNet(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, action_dim),
            nn.Softmax(dim=-1)
        )

    def forward(self, x):
        return self.net(x)