# vision/brainscore_vision/benchmarks/yourbenchmark/__init__.py
from brainscore_vision import benchmark_registry
from .benchmark import YourBenchmarkIT

benchmark_registry['YourBenchmark.IT-pls'] = YourBenchmarkIT

# vision/brainscore_vision/benchmarks/yourbenchmark/benchmark.py
from brainscore_vision import load_dataset, load_metric, load_ceiling
from brainscore_vision.benchmark_helpers.neural_common import NeuralBenchmark


BIBTEX = """@article{YourName2024, ...}"""

def YourBenchmarkIT():
    # Load your registered data
    assembly = load_dataset('YourDataset')
    assembly_repetition = load_dataset('YourDataset')  # Keep repetitions for ceiling

    # Average repetitions for model comparison
    assembly = assembly.mean(dim='repetition')

    return NeuralBenchmark(
        identifier='YourBenchmark.IT-pls',
        version=1,
        assembly=assembly,
        similarity_metric=load_metric('pls'),
        visual_degrees=8,  # From your experiment
        number_of_trials=50,
        ceiling_func=lambda: load_ceiling('internal_consistency')(assembly_repetition),
        parent='IT',  # Or 'V1', 'V2', 'V4'
        bibtex=BIBTEX
    )

┌─────────────────────────────────────────────────────────────────────────────┐
│  Do you have standard neural recordings (electrophysiology/fMRI)?           │
│  ├── YES → Use NeuralBenchmark (Quick Start above)                          │
│  │         Examples: MajajHong2015, FreemanZiemba2013                       │
│  │                                                                          │
│  └── NO → What's special about your data?                                   │
│      │                                                                      │
│      ├── Separate train/test splits?                                        │
│      │   └── Use TrainTestNeuralBenchmark                                   │
│      │       Example: Papale2025                                            │
│      │                                                                      │
│      ├── Temporal dynamics / multiple time bins?                            │
│      │   └── Use BenchmarkBase with custom __call__                         │
│      │       Example: Kar2019                                               │
│      │                                                                      │
│      ├── RSA / representational similarity?                                 │
│      │   └── Use BenchmarkBase with RDM metric                              │
│      │       Example: Coggan2024_fMRI                                       │
│      │                                                                      │
│      └── Neuronal properties (tuning curves, receptive fields)?             │
│          └── Use PropertiesBenchmark                                        │
│              Example: Marques2020                                           │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│  Benchmark (ABC)                                                    │
│  ├── Abstract interface defining required methods                   │
│  ├── __call__(candidate) → Score                                    │
│  ├── identifier, ceiling, version, bibtex properties                │
│  └── Located: brainscore_vision.benchmarks.Benchmark                │
├─────────────────────────────────────────────────────────────────────┤
│  BenchmarkBase(Benchmark)                                           │
│  ├── Helper class implementing standard functions                   │
│  ├── Automatic ceiling caching                                      │
│  ├── Version and metadata management                                │
│  └── Located: brainscore_vision.benchmarks.BenchmarkBase            │
├─────────────────────────────────────────────────────────────────────┤
│  NeuralBenchmark(BenchmarkBase)                                     │
│  ├── Specialized for neural recording comparisons                   │
│  ├── Handles: start_recording(), place_on_screen(), time bins       │
│  ├── Built-in explained_variance ceiling normalization              │
│  └── Located: brainscore_vision.benchmark_helpers.neural_common     │
└─────────────────────────────────────────────────────────────────────┘

from brainscore_vision import load_ceiling

ceiler = load_ceiling('internal_consistency')
benchmark = _DicarloMajajHong2015Region(
    region='IT',
    access='public', 
    identifier_metric_suffix='pls',
    similarity_metric=load_metric('pls'),
    ceiler=ceiler
)

# Located: metrics/regression_correlation/metric.py

class CrossRegressedCorrelation(Metric):
    def __call__(self, source: DataAssembly, target: DataAssembly) -> Score:
        # Cross-validation handles train/test splits
        return self.cross_validation(source, target, apply=self.apply, aggregate=self.aggregate)

    def apply(self, source_train, target_train, source_test, target_test):
        # 1. Fit PLS: model_activations → neural_data (training set)
        self.regression.fit(source_train, target_train)

        # 2. Predict neural data from model activations (test set)
        prediction = self.regression.predict(source_test)

        # 3. Correlate predictions with actual neural data
        score = self.correlation(prediction, target_test)
        return score

    def aggregate(self, scores):
        # Median correlation across neuroids (robust to outliers)
        return scores.median(dim='neuroid')

# Located: metrics/regression_correlation/metric.py

# Standard Ridge (fixed alpha=1)
metric = load_metric('ridge')

# RidgeCV (auto-tunes regularization)
metric = load_metric('ridgecv_split')  # For pre-split train/test data

# Located: metrics/regression_correlation/metric.py

metric = load_metric('neuron_to_neuron')

# Located: metrics/rdm/metric.py

metric = load_metric('rdm')      # Single comparison
metric = load_metric('rdm_cv')   # Cross-validated

# Located: metrics/cka/metric.py

metric = load_metric('cka')      # Single comparison
metric = load_metric('cka_cv')   # Cross-validated

┌─────────────────────────────────────────────────────────────────────────────┐
│  Decision Guide: Which Metric Should I Use?                                 │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  Q: Do you want to predict neural responses from model activations?         │
│  ├── YES → Use regression-based metrics                                     │
│  │   ├── Default: PLS (handles high dimensions well)                        │
│  │   ├── Want regularization control? → Ridge                               │
│  │   └── Want 1:1 unit matching? → Neuron-to-Neuron                         │
│  │                                                                          │
│  └── NO → Use representational similarity metrics                           │
│      ├── Classic RSA paradigm? → RDM                                        │
│      └── Comparing representational geometry? → CKA                         │
│                                                                             │
│  Most Brain-Score neural benchmarks use: pls (default) or ridge             │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

# Cross-validation (most common): metric handles splitting
metric = load_metric('pls')  # Alias for 'pls_cv'

# Fixed split: you provide separate train and test data
metric = load_metric('pls_split')
score = metric(source_train, source_test, target_train, target_test)

metric = load_metric('spantime_pls')    # PLS across time bins
metric = load_metric('spantime_ridge')  # Ridge across time bins

# From MajajHong2015
assembly = load_assembly(average_repetitions=True, region='IT')      # For metric
assembly_repetition = load_assembly(average_repetitions=False, region='IT')  # For ceiling

ceiling_func = lambda: ceiler(assembly_repetition)  # Split-half uses repetitions

# From FreemanZiemba2013
similarity_metric = load_metric('pls', crossvalidation_kwargs=dict(
    stratification_coord='texture_type'  # Balance texture vs. noise images
))

# From Kar2019 - Object Solution Times
TIME_BINS = [(time_bin_start, time_bin_start + 10) 
             for time_bin_start in range(70, 250, 10)]  # 70-250ms in 10ms steps

candidate.start_recording('IT', time_bins=TIME_BINS)

# Static benchmark (default)
TIME_BINS = [(70, 170)]  # Single 100ms window

# Temporal benchmark
TIME_BINS = [(t, t+10) for t in range(70, 250, 10)]  # 18 time bins × 10ms each

# From MajajHong2015 - separate V4 and IT benchmarks
def load_assembly(region: str, average_repetitions: bool):
    assembly = load_dataset('MajajHong2015')
    assembly = assembly.sel(neuroid=assembly['region'] == region)  # Filter by region
    if average_repetitions:
        assembly = assembly.mean(dim='repetition')
    return assembly

# Creates separate benchmarks
benchmark_IT = NeuralBenchmark(identifier='MajajHong2015.IT-pls', ...)
benchmark_V4 = NeuralBenchmark(identifier='MajajHong2015.V4-pls', ...)

# From FreemanZiemba2013 - Texture Modulation properties
def freemanziemba2013_properties(responses, baseline):
    # Uses 'type', 'family', 'sample' coordinates to organize responses
    responses = responses.sortby(['type', 'family', 'sample'])

    type = np.array(sorted(set(responses.type.values)))    # texture vs. noise
    family = np.array(sorted(set(responses.family.values)))  # texture family
    sample = np.array(sorted(set(responses.sample.values)))  # specific sample

    # Reshape using coordinate structure
    responses = responses.values.reshape(n_neuroids, len(type), len(family), len(sample))

    # Compute texture modulation index from structured data
    texture_modulation_index = calc_texture_modulation(responses[:, 1], responses[:, 0])

from brainscore_vision.benchmark_helpers.neural_common import NeuralBenchmark

def DicarloMajajHong2015ITPLS():
    assembly = load_assembly(average_repetitions=True, region='IT')
    assembly_repetition = load_assembly(average_repetitions=False, region='IT')

    return NeuralBenchmark(
        identifier='MajajHong2015.IT-pls',
        assembly=assembly,
        similarity_metric=load_metric('pls'),
        visual_degrees=8,
        number_of_trials=50,
        ceiling_func=lambda: ceiler(assembly_repetition),
        parent='IT',
        bibtex=BIBTEX
    )

from brainscore_vision.benchmark_helpers.properties_common import PropertiesBenchmark

def MarquesCavanaugh2002V1SurroundSuppressionIndex():
    assembly = load_dataset('Cavanaugh2002a')
    similarity_metric = load_metric('ks_similarity', property_name='surround_suppression_index')

    return PropertiesBenchmark(
        identifier='Marques2020_Cavanaugh2002-surround_suppression_index',
        assembly=assembly,
        neuronal_property=cavanaugh2002_properties,
        similarity_metric=similarity_metric,
        timebins=[(70, 170)],
        ceiling_func=NeuronalPropertyCeiling(similarity_metric),
        parent='V1-surround_modulation',
        bibtex=BIBTEX
    )

from brainscore_vision.benchmarks import BenchmarkBase

class _Bracci2019RSA(BenchmarkBase):
    def __init__(self, region):
        self._stimulus_set = load_stimulus_set('Bracci2019')
        self._human_assembly = load_dataset('Bracci2019')
        self._metric = load_metric('rdm')

        super().__init__(
            identifier=f'Bracci2019.{region}-rdm',
            version=1,
            ceiling_func=lambda: 1,
            parent='Bracci2019',
            bibtex=BIBTEX
        )

    def __call__(self, candidate: BrainModel):
        # 1. Start recording
        candidate.start_recording(self._region, [(70, 170)])

        # 2. Scale stimuli
        stimulus_set = place_on_screen(
            self._stimulus_set,
            target_visual_degrees=candidate.visual_degrees(),
            source_visual_degrees=8
        )

        # 3. Get model activations
        dnn_assembly = candidate.look_at(stimulus_set, number_of_trials=1)

        # 4. Custom preprocessing and comparison
        ceiling = self._get_human_ceiling(self._human_assembly)
        similarity = self._metric(dnn_assembly, self._human_assembly)

        score = Score(similarity / ceiling)
        score.attrs['raw'] = similarity
        score.attrs['ceiling'] = ceiling

        return score

from brainscore_vision import load_dataset, load_metric, load_ceiling
from brainscore_vision.benchmark_helpers.neural_common import NeuralBenchmark, average_repetition
from brainscore_core.metrics import Score

BIBTEX = """
@article{author2024,
  title={Your Paper Title},
  author={Author, A.},
  journal={Journal},
  year={2024}
}
"""

# Constants from your experiment
VISUAL_DEGREES = 8      # Stimulus size in degrees of visual angle
NUMBER_OF_TRIALS = 50   # Number of presentations per image


def load_assembly(region: str, average_repetitions: bool = True):
    """Load neural data, optionally averaging across repetitions."""
    assembly = load_dataset('MyExperiment2024')

    # Filter by brain region
    assembly = assembly.sel(neuroid=assembly['region'] == region)

    # Average repetitions for model comparison (keep for ceiling)
    if average_repetitions:
        assembly = average_repetition(assembly)

    return assembly


def MyExperiment2024ITPLS():
    """
    IT cortex benchmark using PLS regression.

    Returns a NeuralBenchmark that compares model activations
    to IT neural responses using partial least squares.
    """
    # Load averaged data for metric computation
    assembly = load_assembly(region='IT', average_repetitions=True)

    # Load non-averaged data for ceiling (needs repetitions for split-half)
    assembly_repetition = load_assembly(region='IT', average_repetitions=False)

    # Internal consistency ceiling
    ceiler = load_ceiling('internal_consistency')

    return NeuralBenchmark(
        identifier='MyExperiment2024.IT-pls',
        version=1,
        assembly=assembly,
        similarity_metric=load_metric('pls'),
        visual_degrees=VISUAL_DEGREES,
        number_of_trials=NUMBER_OF_TRIALS,
        ceiling_func=lambda: ceiler(assembly_repetition),
        parent='IT',
        bibtex=BIBTEX
    )


def MyExperiment2024V4PLS():
    """V4 cortex benchmark using PLS regression."""
    assembly = load_assembly(region='V4', average_repetitions=True)
    assembly_repetition = load_assembly(region='V4', average_repetitions=False)
    ceiler = load_ceiling('internal_consistency')

    return NeuralBenchmark(
        identifier='MyExperiment2024.V4-pls',
        version=1,
        assembly=assembly,
        similarity_metric=load_metric('pls'),
        visual_degrees=VISUAL_DEGREES,
        number_of_trials=NUMBER_OF_TRIALS,
        ceiling_func=lambda: ceiler(assembly_repetition),
        parent='V4',
        bibtex=BIBTEX
    )

# MajajHong2015 - region-based parent (V4 or IT)
NeuralBenchmark(
    identifier=f'MajajHong2015.{region}-pls',
    parent=region,  # 'V4' or 'IT'
    ...
)

# Kar2019 - IT cortex benchmark
super().__init__(
    identifier='Kar2019-ost',
    parent='IT',
    ...
)

# Papale2025 - human fMRI regions
NeuralBenchmark(
    identifier=f'Papale2025.{region}-{metric}',
    parent=region,  # 'V1', 'V2', 'V4', 'IT'
    ...
)

# Coggan2024_fMRI - also region-based
NeuralBenchmark(
    identifier=f'Coggan2024.fMRI.{region}-{metric}',
    parent=region,
    ...
)

# vision/brainscore_vision/benchmarks/majajhong2015/__init__.py

from brainscore_vision import benchmark_registry
from .benchmark import (
    DicarloMajajHong2015V4PLS,
    DicarloMajajHong2015ITPLS,
    MajajHongV4PublicBenchmark,
    MajajHongITPublicBenchmark,
)

# Register each benchmark variant
benchmark_registry['MajajHong2015.V4-pls'] = DicarloMajajHong2015V4PLS
benchmark_registry['MajajHong2015.IT-pls'] = DicarloMajajHong2015ITPLS
benchmark_registry['MajajHong2015.V4.public-pls'] = MajajHongV4PublicBenchmark
benchmark_registry['MajajHong2015.IT.public-pls'] = MajajHongITPublicBenchmark

vision/brainscore_vision/benchmarks/majajhong2015/
├── __init__.py          # Registration (imports from benchmark.py)
├── benchmark.py         # Benchmark implementation (factory functions)
├── test.py              # Unit tests
└── requirements.txt     # Dependencies (optional)

# vision/brainscore_vision/benchmarks/majajhong2015/test.py

import pytest
from brainscore_vision import load_benchmark, load_model

@pytest.mark.parametrize("benchmark, expected", [
    ('MajajHong2015.V4-pls', approx(0.89, abs=0.01)),
    ('MajajHong2015.IT-pls', approx(0.82, abs=0.01)),
])
def test_MajajHong2015(benchmark, expected):
    # Load precomputed features to speed up testing
    filepath = Path(__file__).parent / 'alexnet-majaj2015.private-features.12.nc'
    benchmark = load_benchmark(benchmark)

    # Run with precomputed features instead of full model
    score = run_test(benchmark=benchmark, precomputed_features_filepath=filepath)

    assert score == expected

# Solution: Ensure ceiling_func uses non-averaged data with split-half
ceiling_func=lambda: ceiler(assembly_repetition)  # NOT assembly (averaged)

# Solution: Check stimulus_id alignment
model_ids = set(model_assembly['stimulus_id'].values)
bio_ids = set(biological_assembly['stimulus_id'].values)
assert model_ids == bio_ids, f"Mismatched IDs: {model_ids ^ bio_ids}"

# Solution: Verify the region-to-layer mapping
# Check that model.start_recording() is using correct layer
candidate.start_recording('IT', time_bins=[(70, 170)])
# Ensure your model's layer map includes 'IT' → appropriate layer

# Solution: Ensure assembly has time_bin dimension
assert 'time_bin' in assembly.dims or len(assembly.dims) == 2
# For static benchmarks, NeuralBenchmark squeezes single time bins automatically

# Solution: Check data dimensions
print(f"Samples: {len(assembly['presentation'])}, Features: {len(assembly['neuroid'])}")
# PLS needs samples > features; consider using Ridge for high-dimensional data