class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(20, 64), # quat(x,y,z,w) x 5 = 20
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, 45) # pos(x,y,z) x 15 = 45
        )

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

# 正規化（z-score normalize）関数。公式サンプルより拝借
def normalize(X, axis, savefile=None, epsilon=1e-10):
    Xmean = X.mean(axis=axis)
    Xstd = X.std(axis=axis)

    # Set standard deviation to 1 for very small values to avoid division by near-zero numbers
    Xstd[Xstd < epsilon] = 1

    X = (X - Xmean) / Xstd

    if savefile is not None:
        np.save(savefile + '_mean.npy', Xmean)
        np.save(savefile + '_std.npy', Xstd)

    return X

BATCH_SIZE = 128
LEARNING_RATE = 5e-3
EPOCH = 300

if __name__ == '__main__':

    model_name = 'backpack'
    model_version = 'v1'
    save_dir = os.path.join('models', model_name, model_version)
    os.makedirs(save_dir, exist_ok=True)

    # 訓練データの準備（Mayaから出力したnpyを読み込み）
    data_dir = os.path.join(r'D:\BifrostML\data', model_name) # ← パスは適宜書き換えてください
    input_data = np.load(os.path.join(data_dir, 'input_features.npy'))
    output_data = np.load(os.path.join(data_dir, 'output_features.npy'))

    # 出力データの正規化（z-score normalize）
    normalized_output_data = normalize(output_data, axis=0, savefile=os.path.join(save_dir, 'output'))
    print("output_data min-max:", output_data.min(), output_data.max())
    print("norm output_data min-max:", normalized_output_data.min(), normalized_output_data.max())

    # torch.tensorに変換
    input_tensor = torch.tensor(input_data, dtype=torch.float32).to(device)
    output_tensor = torch.tensor(normalized_output_data, dtype=torch.float32).to(device)

    # 入力と出力をまとめてデータローダーに
    dataset = TensorDataset(input_tensor, output_tensor)
    train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

    # モデル定義
    model = Network().to(device)

    # 損失関数とオプティマイザの定義
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    # 学習
    log_dir = f'runs/{model_name}/{model_version}_batch_{BATCH_SIZE}_lr_{LEARNING_RATE}'
    train(model, train_loader, criterion, optimizer, num_epochs=EPOCH, log_dir=log_dir)

    # 学習したモデルのパラメータを保存
    save_model_params(model, save_dir)

from torch.utils.tensorboard import SummaryWriter

def train(model:nn.Module, train_loader, criterion, optimizer, num_epochs=100, log_dir='runs/exp1'):
    """ 学習しつつLossの推移をTensorBoardに記録する関数 """

    writer = SummaryWriter(log_dir=log_dir) # ← 変更箇所

    model.train()
    for epoch in range(num_epochs):
        total_loss = 0.0
        total_samples  = 0

        for train_x, teacher_y in train_loader:
            train_x = train_x.to(device)
            teacher_y = teacher_y.to(device)

            optimizer.zero_grad()  # 勾配リセット
            pred_y = model(train_x)  # 予測
            loss = criterion(pred_y, teacher_y)  # 損失取得
            loss.backward()  # 勾配計算
            optimizer.step()  # パラメータ更新

            # loss記録
            batch_size = train_x.size(0)
            total_loss += loss.item() * batch_size
            total_samples  += batch_size

        avg_train_loss = total_loss / total_samples
        writer.add_scalar('Loss/train', avg_train_loss, epoch) # ← 変更箇所

        print(f'[Epoch {epoch+1:2d}/{num_epochs:2d}], Train Loss:{avg_train_loss:.4f}')

    writer.close()

# 訓練データの準備（Mayaから出力したnpyを読み込み）
data_dir = os.path.join(r'D:\BifrostML\data', model_name)
input_data = np.load(os.path.join(data_dir, 'input_features.npy'))
output_data = np.load(os.path.join(data_dir, 'output_features.npy'))

# 追加で用意した訓練データの読み込み
input_data_add01 = np.load(os.path.join(data_dir, 'input_features_add01.npy'))
output_data_add01 = np.load(os.path.join(data_dir, 'output_features_add01.npy'))
input_data_add02 = np.load(os.path.join(data_dir, 'input_features_add02.npy'))
output_data_add02 = np.load(os.path.join(data_dir, 'output_features_add02.npy'))

# すべてのデータを結合
#   add01（肩だけ動いているデータ）については、2つおきに10000ペア使用
#   add02（背骨がすべて同じ回転のデータ）については、1つおきに15000ペア使用
input_data = np.concatenate([input_data, input_data_add01[2::3], input_data_add02[1::2]], axis=0)
output_data = np.concatenate([output_data, output_data_add01[2::3], output_data_add02[1::2]], axis=0)

criterion = nn.SmoothL1Loss()

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.fc1 = nn.Linear(20, 128)
        self.fc2 = nn.Linear(128, 256)
        self.fc3 = nn.Linear(256, 256)
        self.fc4 = nn.Linear(256, 45)
        self.activation = nn.Mish()

    def forward(self, x):
        x = self.activation(self.fc1(x))
        res = self.activation(self.fc2(x))
        x = self.fc3(res) + res # Skip Connection
        x = self.activation(x)
        x = self.fc4(x)
        return x