components/TrainingMonitor.jsx

import { useState } from 'react';
import { Card, CardContent, CardHeader, CardTitle } from './ui/Card';
import { Button } from './ui/Button';
import { Play, Pause, Download } from 'lucide-react';
import { LineChart, Line, XAxis, YAxis, CartesianGrid, Tooltip, ResponsiveContainer } from 'recharts';

export default function TrainingMonitor() {
  const [isTraining, setIsTraining] = useState(false);
  const [history, setHistory] = useState({
    loss: [],
    epoch: [],
    lr: [],
  });

  const startTraining = async () => {
    setIsTraining(true);
    try {
      const response = await fetch('/api/training/start', {
        method: 'POST',
      });
      
      const data = await response.json();
      setHistory(data.history);
    } catch (error) {
      console.error('Training failed:', error);
    } finally {
      setIsTraining(false);
    }
  };

  const stopTraining = () => {
    setIsTraining(false);
  };

  const chartData = history.epoch.map((epoch, idx) => ({
    epoch,
    loss: history.loss[idx],
    lr: history.lr[idx],
  }));

  return (
    <div className="space-y-6">
      <h2 className="text-3xl font-bold text-gray-800">Training & Submit</h2>
      
      <Card>
        <CardHeader>
          <CardTitle>Training Monitor</CardTitle>
        </CardHeader>
        <CardContent>
          <div className="flex gap-4 mb-6">
            <Button onClick={startTraining} disabled={isTraining}>
              <Play className="w-4 h-4 mr-2" />
              Start Training
            </Button>
            <Button onClick={stopTraining} variant="secondary" disabled={!isTraining}>
              <Pause className="w-4 h-4 mr-2" />
              Stop
            </Button>
            <Button variant="outline">
              <Download className="w-4 h-4 mr-2" />
              Export Model
            </Button>
          </div>

          {isTraining && (
            <div className="mb-6">
              <div className="flex items-center justify-center py-8">
                <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-primary-500"></div>
              </div>
              <p className="text-center text-gray-600">Training in progress...</p>
            </div>
          )}

          {chartData.length > 0 && (
            <div className="space-y-6">
              <div>
                <h3 className="text-lg font-semibold mb-4">Training Loss</h3>
                <ResponsiveContainer width="100%" height={300}>
                  <LineChart data={chartData}>
                    <CartesianGrid strokeDasharray="3 3" />
                    <XAxis dataKey="epoch" />
                    <YAxis />
                    <Tooltip />
                    <Line type="monotone" dataKey="loss" stroke="#FF4B4B" />
                  </LineChart>
                </ResponsiveContainer>
              </div>

              <div>
                <h3 className="text-lg font-semibold mb-4">Learning Rate</h3>
                <ResponsiveContainer width="100%" height={200}>
                  <LineChart data={chartData}>
                    <CartesianGrid strokeDasharray="3 3" />
                    <XAxis dataKey="epoch" />
                    <YAxis />
                    <Tooltip />
                    <Line type="monotone" dataKey="lr" stroke="#764ba2" />
                  </LineChart>
                </ResponsiveContainer>
              </div>
            </div>
          )}
        </CardContent>
      </Card>
    </div>
  );
}