frontend/src/pages/AddModelPage/components/ModelSubmissionForm/ModelSubmissionForm.js

import React, { useState } from "react";
import {
  Box,
  Paper,
  Typography,
  TextField,
  Button,
  FormControl,
  FormLabel,
  RadioGroup,
  FormControlLabel,
  Radio,
  InputLabel,
  Select,
  MenuItem,
  Stack,
  Grid,
  CircularProgress,
  Alert,
} from "@mui/material";
import RocketLaunchIcon from "@mui/icons-material/RocketLaunch";
import CheckCircleOutlineIcon from "@mui/icons-material/CheckCircleOutline";
import { alpha } from "@mui/material/styles";
import InfoIconWithTooltip from "../../../../components/shared/InfoIconWithTooltip";
import {
  FINETUNING_STRATEGIES,
  HEADS,
  PEFT_STRATEGIES,
  OEB_DATASETS,
} from "../../../../pages/LeaderboardPage/components/Leaderboard/constants/modelTypes";
import AuthContainer from "../../../../components/shared/AuthContainer";

const HELP_TEXTS = {
  modelClass: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Model Class (dotted import path)
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Fully-qualified import path to the model class evaluated by
        OpenEEGBench. The class must accept (batch, n_chans, n_times) and
        return (batch, n_outputs). Example: braindecode.models.BIOT
      </Typography>
    </Box>
  ),
  modelName: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Model Name
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        A human-readable name for your submission. This is how your model will
        appear on the leaderboard.
      </Typography>
    </Box>
  ),
  weightsSource: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Weights Source
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Provide exactly one source for the pre-trained weights: a public
        Hugging Face Hub repository (with a license tag) or a direct
        http(s) URL to a checkpoint.
      </Typography>
    </Box>
  ),
  headModuleName: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Head Module Name
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Name of the model's classification head module. OpenEEGBench replaces
        this module to attach evaluation heads. Default: final_layer
      </Typography>
    </Box>
  ),
  finetuningStrategies: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Fine-tuning Strategies
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Frozen: linear probing (encoder frozen). LoRA / IA3 / AdaLoRA / DoRA /
        OFT: parameter-efficient fine-tuning. Full fine-tune: all parameters
        updated. Two stages: probe then fine-tune. Select one or more.
      </Typography>
    </Box>
  ),
  heads: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Evaluation Heads
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Linear head: single linear layer. MLP head: small multi-layer
        perceptron. Original head: keep the model's own head. Select one or
        more.
      </Typography>
    </Box>
  ),
  peftTargetModules: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        PEFT Target Modules
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Comma-separated module names to which the PEFT adapter is applied
        (e.g. to_q, to_v). Required when a PEFT strategy (LoRA, IA3, AdaLoRA,
        DoRA, OFT) is selected.
      </Typography>
    </Box>
  ),
  peftFfModules: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        PEFT Feed-Forward Modules
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Comma-separated feed-forward module names scaled by IA3. Required when
        the IA3 strategy is selected.
      </Typography>
    </Box>
  ),
  datasets: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Datasets
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        The downstream datasets to evaluate on. Leave empty to run all 10
        datasets in the OpenEEGBench classification suite.
      </Typography>
    </Box>
  ),
  nSeeds: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Number of Seeds
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        How many random seeds to average over per dataset. More seeds give a
        more stable estimate but take longer. Default: 3
      </Typography>
    </Box>
  ),
  contactEmail: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Contact Email
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Optional. Stored privately and used only to reach you about your
        submission. It is never shown on the public leaderboard.
      </Typography>
    </Box>
  ),
  visibility: (
    <Box sx={{ p: 1 }}>
      <Typography variant="subtitle2" sx={{ fontWeight: 600, mb: 0.5 }}>
        Visibility
      </Typography>
      <Typography variant="body2" sx={{ opacity: 0.9, lineHeight: 1.4 }}>
        Public submissions appear on the leaderboard. Private submissions are
        evaluated but kept out of the public table.
      </Typography>
    </Box>
  ),
};

const VISIBILITY_OPTIONS = ["Public", "Private"];

const INITIAL_FORM_DATA = {
  modelClass: "", // dotted import path
  modelName: "",
  organization: "",
  weightsSource: "hub_repo", // "hub_repo" | "checkpoint_url"
  hubRepo: "",
  checkpointUrl: "",
  headModuleName: "final_layer",
  finetuningStrategies: ["frozen"],
  heads: ["linear_head"],
  peftTargetModules: "", // comma-separated -> array on submit
  peftFfModules: "", // comma-separated -> array on submit
  datasets: [], // empty = all
  nSeeds: 3,
  modelUrl: "",
  paperUrl: "",
  contactEmail: "",
  additionalInfo: "",
  visibility: "Public",
};

const toList = (s) =>
  s
    .split(",")
    .map((x) => x.trim())
    .filter(Boolean);

const strategyLabel = (value) =>
  FINETUNING_STRATEGIES.find((s) => s.value === value)?.label || value;

const headLabel = (value) =>
  HEADS.find((h) => h.value === value)?.label || value;

function ModelSubmissionForm({ user, isAuthenticated }) {
  const [formData, setFormData] = useState(INITIAL_FORM_DATA);
  const [error, setError] = useState(null);
  const [submitting, setSubmitting] = useState(false);
  const [success, setSuccess] = useState(false);
  const [submittedData, setSubmittedData] = useState(null);

  const showPeftTargetModules = formData.finetuningStrategies.some((s) =>
    PEFT_STRATEGIES.includes(s)
  );
  const showPeftFfModules = formData.finetuningStrategies.includes("ia3");

  const handleChange = (event) => {
    const { name, value, checked } = event.target;
    setFormData((prev) => ({
      ...prev,
      [name]: event.target.type === "checkbox" ? checked : value,
    }));
  };

  const handleSubmit = async (e) => {
    e.preventDefault();
    setError(null);
    setSubmitting(true);

    try {
      const body = {
        user_id: user.username,
        model_cls: formData.modelClass,
        model_name: formData.modelName,
        organization: formData.organization,
        hub_repo:
          formData.weightsSource === "hub_repo" ? formData.hubRepo : "",
        checkpoint_url:
          formData.weightsSource === "checkpoint_url"
            ? formData.checkpointUrl
            : "",
        head_module_name: formData.headModuleName,
        finetuning_strategies: formData.finetuningStrategies,
        heads: formData.heads,
        peft_target_modules: toList(formData.peftTargetModules),
        peft_ff_modules: toList(formData.peftFfModules),
        datasets: formData.datasets,
        n_seeds: Number(formData.nSeeds),
        model_url: formData.modelUrl,
        paper_url: formData.paperUrl,
        contact_email: formData.contactEmail,
        additional_info: formData.additionalInfo,
        visibility: formData.visibility,
      };

      const response = await fetch("/api/models/submit", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify(body),
      });

      if (!response.ok) {
        const error = await response.json();
        throw new Error(error.detail || "Failed to submit model");
      }

      setSubmittedData(formData);
      setSuccess(true);
    } catch (error) {
      setError(error.message);
    } finally {
      setSubmitting(false);
    }
  };

  if (success && submittedData) {
    return (
      <Paper
        variant="outlined"
        sx={(theme) => ({
          p: 6,
          mb: 3,
          bgcolor: alpha(theme.palette.success.main, 0.05),
          borderColor: alpha(theme.palette.success.main, 0.2),
        })}
      >
        <Stack spacing={3}>
          <Stack direction="row" spacing={2} alignItems="center">
            <CheckCircleOutlineIcon color="success" sx={{ fontSize: 28 }} />
            <Typography
              variant="h5"
              sx={{ fontWeight: 600, color: "success.800" }}
            >
              Model submitted successfully!
            </Typography>
          </Stack>

          <Typography variant="body1">
            Your model <strong>{submittedData.modelName}</strong> has been added
            to the evaluation queue with the following parameters:
          </Typography>

          <Paper
            variant="outlined"
            sx={{
              p: 2,
              borderColor: "divider",
            }}
          >
            <Stack spacing={1.5}>
              <Stack direction="row" spacing={2}>
                <Typography
                  variant="body2"
                  color="text.secondary"
                  sx={{ width: 120 }}
                >
                  Model:
                </Typography>
                <Typography variant="body2" sx={{ fontFamily: "monospace" }}>
                  {submittedData.modelName}
                </Typography>
              </Stack>
              <Stack direction="row" spacing={2}>
                <Typography
                  variant="body2"
                  color="text.secondary"
                  sx={{ width: 120 }}
                >
                  Model class:
                </Typography>
                <Typography variant="body2" sx={{ fontFamily: "monospace" }}>
                  {submittedData.modelClass}
                </Typography>
              </Stack>
              <Stack direction="row" spacing={2}>
                <Typography
                  variant="body2"
                  color="text.secondary"
                  sx={{ width: 120 }}
                >
                  Strategies:
                </Typography>
                <Typography variant="body2">
                  {submittedData.finetuningStrategies
                    .map(strategyLabel)
                    .join(", ")}
                </Typography>
              </Stack>
              <Stack direction="row" spacing={2}>
                <Typography
                  variant="body2"
                  color="text.secondary"
                  sx={{ width: 120 }}
                >
                  Visibility:
                </Typography>
                <Typography variant="body2">
                  {submittedData.visibility}
                </Typography>
              </Stack>
            </Stack>
          </Paper>

          <Typography variant="body2" color="text.secondary">
            A worker will run oeb.benchmark() and post the official scores to
            the leaderboard.
          </Typography>

          <Stack direction="row" spacing={2}>
            <Button
              variant="outlined"
              size="large"
              onClick={() => {
                setSuccess(false);
                setSubmittedData(null);
                setFormData(INITIAL_FORM_DATA);
              }}
            >
              Submit another model
            </Button>
          </Stack>
        </Stack>
      </Paper>
    );
  }

  return (
    <>
      {error && (
        <Alert severity="error" sx={{ mb: 2 }}>
          {error}
        </Alert>
      )}
      <AuthContainer actionText="submit a model" />
      {isAuthenticated && (
        <Paper
          elevation={0}
          component="form"
          onSubmit={handleSubmit}
          sx={{
            p: 0,
            border: "1px solid",
            borderColor: "grey.300",
            mb: 3,
            overflow: "hidden",
          }}
        >
          {/* Header */}
          <Box
            sx={{
              px: 3,
              py: 2,
              borderBottom: "1px solid",
              borderColor: (theme) =>
                theme.palette.mode === "dark"
                  ? alpha(theme.palette.divider, 0.1)
                  : "grey.200",
              bgcolor: (theme) =>
                theme.palette.mode === "dark"
                  ? alpha(theme.palette.background.paper, 0.5)
                  : "grey.50",
            }}
          >
            <Typography
              variant="h6"
              sx={{ fontWeight: 600, color: "text.primary" }}
            >
              Model Submission Form
            </Typography>
          </Box>

          {/* Form Content */}
          <Box sx={{ p: 3 }}>
            <Grid container spacing={3}>
              {/* Model Information */}
              <Grid item xs={12}>
                <Stack direction="row" spacing={1} alignItems="center">
                  <Typography variant="h6">Model Information</Typography>
                  <InfoIconWithTooltip tooltip={HELP_TEXTS.modelClass} />
                </Stack>
              </Grid>

              <Grid item xs={12}>
                <TextField
                  required
                  fullWidth
                  name="modelClass"
                  label="Model Class"
                  placeholder="braindecode.models.BIOT"
                  value={formData.modelClass}
                  onChange={handleChange}
                  helperText="e.g. braindecode.models.BIOT"
                  InputProps={{
                    endAdornment: (
                      <InfoIconWithTooltip tooltip={HELP_TEXTS.modelClass} />
                    ),
                  }}
                />
              </Grid>

              <Grid item xs={12} sm={6}>
                <TextField
                  required
                  fullWidth
                  name="modelName"
                  label="Model Name"
                  placeholder="BIOT-frozen"
                  value={formData.modelName}
                  onChange={handleChange}
                  helperText="Display name on the leaderboard"
                  InputProps={{
                    endAdornment: (
                      <InfoIconWithTooltip tooltip={HELP_TEXTS.modelName} />
                    ),
                  }}
                />
              </Grid>

              <Grid item xs={12} sm={6}>
                <TextField
                  required
                  fullWidth
                  name="organization"
                  label="Organization"
                  placeholder="braindecode"
                  value={formData.organization}
                  onChange={handleChange}
                  helperText="Author or organization name"
                />
              </Grid>

              {/* Weights Source */}
              <Grid item xs={12}>
                <Stack direction="row" spacing={1} alignItems="center">
                  <Typography variant="h6">Weights Source</Typography>
                  <InfoIconWithTooltip tooltip={HELP_TEXTS.weightsSource} />
                </Stack>
              </Grid>

              <Grid item xs={12}>
                <FormControl>
                  <FormLabel id="weights-source-label">
                    Provide one weights source
                  </FormLabel>
                  <RadioGroup
                    row
                    aria-labelledby="weights-source-label"
                    name="weightsSource"
                    value={formData.weightsSource}
                    onChange={handleChange}
                  >
                    <FormControlLabel
                      value="hub_repo"
                      control={<Radio />}
                      label="Hugging Face Hub repo"
                    />
                    <FormControlLabel
                      value="checkpoint_url"
                      control={<Radio />}
                      label="Checkpoint URL"
                    />
                  </RadioGroup>
                </FormControl>
              </Grid>

              {formData.weightsSource === "hub_repo" ? (
                <Grid item xs={12}>
                  <TextField
                    fullWidth
                    name="hubRepo"
                    label="Hub Repository"
                    placeholder="organization/model-weights"
                    value={formData.hubRepo}
                    onChange={handleChange}
                    helperText="Public HF Hub repo containing the weights (with a license tag)"
                  />
                </Grid>
              ) : (
                <Grid item xs={12}>
                  <TextField
                    fullWidth
                    name="checkpointUrl"
                    label="Checkpoint URL"
                    placeholder="https://example.com/weights.pth"
                    value={formData.checkpointUrl}
                    onChange={handleChange}
                    helperText="Direct http(s) URL to a checkpoint file"
                  />
                </Grid>
              )}

              <Grid item xs={12}>
                <TextField
                  fullWidth
                  name="headModuleName"
                  label="Head Module Name"
                  value={formData.headModuleName}
                  onChange={handleChange}
                  helperText="Default: final_layer"
                  InputProps={{
                    endAdornment: (
                      <InfoIconWithTooltip
                        tooltip={HELP_TEXTS.headModuleName}
                      />
                    ),
                  }}
                />
              </Grid>

              {/* Evaluation Configuration */}
              <Grid item xs={12}>
                <Stack direction="row" spacing={1} alignItems="center">
                  <Typography variant="h6">Evaluation Configuration</Typography>
                </Stack>
              </Grid>

              <Grid item xs={12} sm={6}>
                <FormControl fullWidth>
                  <InputLabel>Fine-tuning Strategies</InputLabel>
                  <Select
                    multiple
                    name="finetuningStrategies"
                    value={formData.finetuningStrategies}
                    onChange={handleChange}
                    label="Fine-tuning Strategies"
                    renderValue={(selected) =>
                      selected.map(strategyLabel).join(", ")
                    }
                    endAdornment={
                      <InfoIconWithTooltip
                        tooltip={HELP_TEXTS.finetuningStrategies}
                        sx={{ mr: 2 }}
                      />
                    }
                  >
                    {FINETUNING_STRATEGIES.map((strategy) => (
                      <MenuItem key={strategy.value} value={strategy.value}>
                        {strategy.label}
                      </MenuItem>
                    ))}
                  </Select>
                </FormControl>
              </Grid>

              <Grid item xs={12} sm={6}>
                <FormControl fullWidth>
                  <InputLabel>Heads</InputLabel>
                  <Select
                    multiple
                    name="heads"
                    value={formData.heads}
                    onChange={handleChange}
                    label="Heads"
                    renderValue={(selected) =>
                      selected.map(headLabel).join(", ")
                    }
                    endAdornment={
                      <InfoIconWithTooltip
                        tooltip={HELP_TEXTS.heads}
                        sx={{ mr: 2 }}
                      />
                    }
                  >
                    {HEADS.map((head) => (
                      <MenuItem key={head.value} value={head.value}>
                        {head.label}
                      </MenuItem>
                    ))}
                  </Select>
                </FormControl>
              </Grid>

              {showPeftTargetModules && (
                <Grid item xs={12} sm={showPeftFfModules ? 6 : 12}>
                  <TextField
                    fullWidth
                    name="peftTargetModules"
                    label="PEFT Target Modules"
                    placeholder="to_q, to_v"
                    value={formData.peftTargetModules}
                    onChange={handleChange}
                    helperText="comma-separated module names"
                    InputProps={{
                      endAdornment: (
                        <InfoIconWithTooltip
                          tooltip={HELP_TEXTS.peftTargetModules}
                        />
                      ),
                    }}
                  />
                </Grid>
              )}

              {showPeftFfModules && (
                <Grid item xs={12} sm={6}>
                  <TextField
                    fullWidth
                    name="peftFfModules"
                    label="PEFT Feed-Forward Modules"
                    placeholder="to_ff"
                    value={formData.peftFfModules}
                    onChange={handleChange}
                    helperText="comma-separated module names (IA3 only)"
                    InputProps={{
                      endAdornment: (
                        <InfoIconWithTooltip
                          tooltip={HELP_TEXTS.peftFfModules}
                        />
                      ),
                    }}
                  />
                </Grid>
              )}

              <Grid item xs={12} sm={8}>
                <FormControl fullWidth>
                  <InputLabel>Datasets</InputLabel>
                  <Select
                    multiple
                    name="datasets"
                    value={formData.datasets}
                    onChange={handleChange}
                    label="Datasets"
                    renderValue={(selected) =>
                      selected.length ? selected.join(", ") : "All 10 datasets"
                    }
                    endAdornment={
                      <InfoIconWithTooltip
                        tooltip={HELP_TEXTS.datasets}
                        sx={{ mr: 2 }}
                      />
                    }
                  >
                    {OEB_DATASETS.map((dataset) => (
                      <MenuItem key={dataset} value={dataset}>
                        {dataset}
                      </MenuItem>
                    ))}
                  </Select>
                </FormControl>
              </Grid>

              <Grid item xs={12} sm={4}>
                <TextField
                  fullWidth
                  type="number"
                  name="nSeeds"
                  label="Number of Seeds"
                  value={formData.nSeeds}
                  onChange={handleChange}
                  helperText="Default: 3"
                  inputProps={{ min: 1 }}
                  InputProps={{
                    endAdornment: (
                      <InfoIconWithTooltip tooltip={HELP_TEXTS.nSeeds} />
                    ),
                  }}
                />
              </Grid>

              {/* Additional Information */}
              <Grid item xs={12}>
                <Stack direction="row" spacing={1} alignItems="center">
                  <Typography variant="h6">Additional Information</Typography>
                </Stack>
              </Grid>

              <Grid item xs={12} sm={6}>
                <TextField
                  fullWidth
                  name="modelUrl"
                  label="Model URL"
                  placeholder="https://huggingface.co/organization/model"
                  value={formData.modelUrl}
                  onChange={handleChange}
                  helperText="Optional. Link to the model card or repo."
                />
              </Grid>

              <Grid item xs={12} sm={6}>
                <TextField
                  fullWidth
                  name="paperUrl"
                  label="Paper URL"
                  placeholder="https://arxiv.org/abs/..."
                  value={formData.paperUrl}
                  onChange={handleChange}
                  helperText="Optional. Link to the related paper."
                />
              </Grid>

              <Grid item xs={12} sm={6}>
                <TextField
                  fullWidth
                  name="contactEmail"
                  label="Contact Email"
                  placeholder="you@example.com"
                  value={formData.contactEmail}
                  onChange={handleChange}
                  helperText="Optional. Stored privately."
                  InputProps={{
                    endAdornment: (
                      <InfoIconWithTooltip tooltip={HELP_TEXTS.contactEmail} />
                    ),
                  }}
                />
              </Grid>

              <Grid item xs={12} sm={6}>
                <FormControl fullWidth>
                  <InputLabel>Visibility</InputLabel>
                  <Select
                    name="visibility"
                    value={formData.visibility}
                    onChange={handleChange}
                    label="Visibility"
                    endAdornment={
                      <InfoIconWithTooltip
                        tooltip={HELP_TEXTS.visibility}
                        sx={{ mr: 2 }}
                      />
                    }
                  >
                    {VISIBILITY_OPTIONS.map((option) => (
                      <MenuItem key={option} value={option}>
                        {option}
                      </MenuItem>
                    ))}
                  </Select>
                </FormControl>
              </Grid>

              <Grid item xs={12}>
                <TextField
                  fullWidth
                  multiline
                  minRows={2}
                  name="additionalInfo"
                  label="Additional Info"
                  value={formData.additionalInfo}
                  onChange={handleChange}
                  helperText="Optional. Anything else we should know."
                />
              </Grid>

              {/* Submit Button */}
              <Grid item xs={12}>
                <Box
                  sx={{
                    display: "flex",
                    justifyContent: "space-between",
                    alignItems: "center",
                    mt: 2,
                  }}
                >
                  <Typography variant="body2" color="text.secondary">
                    All fields marked with * are required
                  </Typography>
                  <Button
                    type="submit"
                    variant="contained"
                    disabled={submitting}
                    endIcon={submitting ? null : <RocketLaunchIcon />}
                    sx={{
                      minWidth: 120,
                      position: "relative",
                    }}
                  >
                    {submitting ? (
                      <CircularProgress size={24} color="inherit" />
                    ) : (
                      "Submit"
                    )}
                  </Button>
                </Box>
              </Grid>
            </Grid>
          </Box>
        </Paper>
      )}
    </>
  );
}

export default ModelSubmissionForm;