#!/bin/bash

## Prepare data
mkdir -p data/src
mkdir -p data/processed
mkdir tmp
# FB15K-237
git clone git@github.com:villmow/datasets_knowledge_embedding.git tmp/datasets
mv tmp/datasets/FB15k-237 data/src/FB15K-237
rm -rf tmp/datasets
uv run scripts/preprocess.py --datasets FB15K-237 --data-folder data/src --output-folder data/processed
# Hetionet, openbiolink
uv run data/download_hetionet.py --output_dir data/processed/Hetionet --seed 42
uv run data/download_openbiolink.py --output_dir data/processed/openbiolink

## Run experiments
datasets=("FB15K-237" "WN18RR" "Hetionet" "ogbl-biokg" "openbiolink")
dimensions=(200 1000)

mkdir -p scripts/experiments

for dataset in "${datasets[@]}"; do
    for dim in "${dimensions[@]}"; do
        # ConvE experiment
        run_name="${dataset}_tail_${dim}_conve"
        script_path="scripts/experiments/${run_name}.sh"
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_model dimension=$dim model/fusing_function=conve model.fusing_dropout=0.0 engine_config.learning_rate=1e-4" >> "$script_path"
        chmod +x "$script_path"

        # Hadamard experiment
        run_name="${dataset}_tail_${dim}_hadamard"
        script_path="scripts/experiments/${run_name}.sh"
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_model dimension=$dim model/fusing_function=hadamard model.fusing_dropout=0.1 engine_config.learning_rate=1e-3" >> "$script_path"
        chmod +x "$script_path"

        # RESCAL experiment
        run_name="${dataset}_tail_${dim}_rescal"
        script_path="scripts/experiments/${run_name}.sh"
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_model dimension=$dim model/embedding=rescal model/fusing_function=rescal model.fusing_dropout=0.1 engine_config.learning_rate=1e-4" >> "$script_path"
        chmod +x "$script_path"

        # ComplEx experiment
        run_name="${dataset}_tail_${dim}_complex"
        script_path="scripts/experiments/${run_name}.sh"
        actual_dim=$((dim / 2)) # complex uses half the dimension
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_model dimension=$actual_dim model/embedding=complex model/fusing_function=complex model.fusing_dropout=0.1 engine_config.learning_rate=1e-3" >> "$script_path"
        chmod +x "$script_path"

        # Mixture ConvE experiment
        run_name="${dataset}_mixture_${dim}_conve"
        script_path="scripts/experiments/${run_name}.sh"
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_mixture dimension=$dim model/fusing_function=conve model.mixture_layer.num_components=4 model.fusing_dropout=0.0 engine_config.learning_rate=1e-4" >> "$script_path"
        chmod +x "$script_path"

        # Mixture Hadamard experiment
        run_name="${dataset}_mixture_${dim}_hadamard"
        script_path="scripts/experiments/${run_name}.sh"
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_mixture dimension=$dim model/fusing_function=hadamard model.mixture_layer.num_components=4 model.fusing_dropout=0.1 engine_config.learning_rate=1e-4" >> "$script_path"
        chmod +x "$script_path"

        # Mixture ComplEx experiment
        run_name="${dataset}_mixture_${dim}_complex"
        script_path="scripts/experiments/${run_name}.sh"
        actual_dim=$((dim / 2)) # complex uses half the dimension
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_mixture dimension=$actual_dim model/embedding=complex model/fusing_function=complex model.mixture_layer.num_components=4 model.fusing_dropout=0.1 engine_config.learning_rate=1e-4" >> "$script_path"
        chmod +x "$script_path"

        # Mixture RESCAL experiment
        run_name="${dataset}_mixture_${dim}_rescal"
        script_path="scripts/experiments/${run_name}.sh"
        echo '#!/bin/bash' > "$script_path"
        echo "uv run scripts/main.py dataset=$dataset model=tail_mixture dimension=$dim model.mixture_layer.num_components=4 model/embedding=rescal model/fusing_function=rescal model.fusing_dropout=0.1 engine_config.learning_rate=1e-4" >> "$script_path"
        chmod +x "$script_path"
    done
done

echo "Created experiment scripts in scripts/experiments/"


