export LANG=ja_JP.UTF-8
export LC_ALL=ja_JP.UTF-8

display_custom_help() {
    echo "----------------------------------------------------------------------------------------------------------------------"
    printf "%s\n" "$(conda env list)"
    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "LLMs"
    echo "---"
    echo "conda activate openwebui && open-webui serve --port 6969"
    echo "ollama serve"
    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "Taggers + Captioners"
    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "JTP2"
    echo "---"
    echo "~/toolkit/jtp2 <dir>"
    echo "Joy Captioner"
    echo "---"
    echo "~/source/repos/joy/joy <dir> --custom_prompt \"<prompt>\" --caption_type custom"
    echo "Waifu Diffusion Tagger:"
    echo "---"
    echo "python ~/source/repos/wdv3-timm/wdv3_timm.py <dir> --model eva02"
    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "Database Stuff"
    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "Redis"
    echo "---"
    echo "~/db/redis-stable/src/redis-server      : Start server."
    echo "PostgreSQL"
    echo "---"
    echo "psql -d postgres -h /tmp                : Connect using socket directory."
    echo "Start server:"
    echo "pg_ctl -D \$HOME/db/postgresql/data -l \$HOME/db/pgsql.log start"
#    echo "Commands, Aliases, and Custom Functions:"
#    echo "----------------------------------------------------------------------------------------------------------------------"
#    echo "pie                                         : \`pip install -e . --use-pep517\`"
#    echo "gcs                                         : \`git clone --recurse-submodules\`"
#    echo "dust                                        : A more intuitive version of du."
#    echo "ranger                                      : A vim inspired file manager."
#    echo "htop                                        : Interactive process viewer."
#    echo "nvtop                                       : Interactive GPU process viewer."
#    echo "nvitop                                      : An even more interactive GPU process viewer."
#    echo "nvim                                        : Alias for vim."
#    echo "rt                                          : Edit tmux config and reload it."
#    echo "zr                                          : Edit zsh config and reload it."
#    echo "ta                                          : Attach to tmux session."
#    echo "ga                                          : Git add, commit, and push."
#    echo "gs                                          : Git status."
#    echo "wd                                          : Word diff in git."
#    echo "grabber                                     : Alias for Grabber-cli."
#    echo "ls                                          : Alias for 'ls --color=always'."
#    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "- 🐺 TOOLS                                                                                                           -"
    echo "----------------------------------------------------------------------------------------------------------------------"
    echo "nv                                          : Returns the cuda version number."iexport LANG=ja_JP.UTF-8
export LC_ALL=ja_JP.UTF-8


    echo "remove_repetition                           : Removes repetition in txt files in a target directory."
    echo "copy_sample_prompts                         : Copies ./sample-prompt.txt file from the current dir to datasets/furry."
    echo "remove_number_prefix                        : Removes all numbers prefixed by a _ from the end of every file."
    echo "count_captions                              : Counts *.caption and *.txt files in each subdirectory."
    echo "count_captions_per_folder                   : Counts *.caption and *.txt files in each subdirectory individually."
    echo "llama                                       : Runs Meta-Llama-3-8B-Instruct on port 6969."
    echo "copy_matching_caption_files                 : Copies matching .caption files for <dir> to the current directory."
    echo "c                                           : Change to ComfyUI directory and start the server."
    echo "t                                           : Start TensorBoard with logs directory."
    echo "png2mp4                                     : Convert PNG sequence to MP4 video."
    echo "seed <file>                                 : Display the seed from a safetensors file."
    echo "swch <branch>                               : Clean repo and switch to specified git branch."
    echo "convert_to_jxl <directory>                  : Convert JPG, JPEG, and PNG files to JXL in the specified directory."
    echo "convert_pxl_to_png <directory>              : Convert PXL files to PNG in the specified directory."
    echo "replace_text_in_files [dir] <src> <replace> : Perform text replacement on *.txt files in a target directory."
    echo "update_dir [directory]                      : Update git repositories in subdirectories."
    echo "inject_to_captions [dir] \"txt\"              : Add prefix to the beginning of each text file in a directory."
    echo "chop_lora <input_file>                      : Generate multiple versions of a Lora file with different presets."
    echo "----------------------------------------------------------------------------------------------------------------------"
}

export RUST_BACKTRACE=1

# This function `nv` retrieves the version of the NVIDIA CUDA Compiler (nvcc) installed on the system.
# It extracts the version number from the `nvcc --version` command output.
# The version number is then formatted by removing the dot (e.g., 12.6 becomes 126).
# Finally, the function returns the formatted version number.
function nv() {
    # Get the nvcc version output
    local nvcc_output=$(nvcc --version)

    # Extract the version number (12.6)
    local version=$(echo "$nvcc_output" | grep -oP 'release \K[0-9]+\.[0-9]+')

    # Remove the dot to get 126
    local result=$(echo "$version" | tr -d '.')

    # Print the result
    echo $result
}

export BNB_CUDA_VERSION=126

# Function to remove consecutive repeated words in text files within a directory
remove_repetition() {
  local dir=$1  # The directory to search for text files
  # Find all .txt files in the specified directory and process each file
  find "$dir" -type f -name "*.txt" | while read -r file; do
    # Use awk to process each line of the file
    awk '
    {
      n = split($0, words, " ")  # Split the line into words
      for (i = n; i > 1; i--) {  # Iterate from the last word to the second word
        if (words[i] != words[i-1]) break  # Stop if the current word is not equal to the previous word
      }
      for (j = 1; j <= i; j++) {  # Print the words up to the point where repetition ends
        printf "%s%s", words[j], (j == i ? ORS : OFS)  # Print the word followed by a space or newline
      }
    }
    ' "$file" > "${file}.tmp" && mv "${file}.tmp" "$file"  # Save the processed content to a temporary file and replace the original file
  done
}

# This alias 'pie' is a shortcut for installing a Python package in editable mode
# using the pip command with the --use-pep517 option.
alias pie='pip install -e . --use-pep517'

# Function to remove specific tags from all *.txt files in a target directory recursively
remove_boys() {
    # Assign the first argument passed to the function to the variable target_dir
    local target_dir="$1"
    
    # Find all *.txt files in the target directory and its subdirectories
    find "$target_dir" -type f -name "*.txt" | while read -r file; do
        # Use sed to remove occurrences of [1-9]boy, [1-9]boys, [1-9]girl, and [1-9]girls along with a comma and space character
        # -i.bak creates a backup of the original file with a .bak extension
        # -E enables extended regular expressions
        sed -i.bak -E 's/, ([1-9]boy|[1-9]boys|[1-9]girl|[1-9]girls)//g' "$file"
    done
}

export DOTNET_CLI_TELEMETRY_OPTOUT=1

# Organizes a sample prompt file from the current directory to datasets/furry.
# It moves the file named sample-prompts.txt to either
# ~/datasets/furry/sample_prompts/pony or ~/datasets/furry/sample_prompts/compass based on the content.
# If the file contains the regexp 'score_*', it is moved to ~/datasets/furry/sample_prompts/pony.
# Otherwise, it is moved to ~/datasets/furry/sample_prompts/compass.
# The -v option is used with cp to provide verbose output.
copy_sample_prompts() {
    file="./sample-prompts.txt"
    if grep -q 'score_*' "$file"; then
        cp -v "$file" ~/datasets/furry/sample_prompts/pony/
    else
        cp -v "$file" ~/datasets/furry/sample_prompts/compass/
    fi

    echo "File has been organized."
}

# Removes all numbers prefixed by a _ from the end of every file before the file extension
remove_number_prefix() {
  # Loop through all files in the current directory and its subdirectories
  for file in **/*_[0-9]*.*; do
    # Get the new file name by removing '_number' before the file extension
    new_file="${file%_[0-9]*.*}.${file##*.}"
    # Rename the file
    mv "$file" "$new_file"
  done
}

# Counts all *.caption and *.txt files in all subdirectories.
count_captions() {
  caption_count=$(find . -type f -name "*.caption" | wc -l)
  txt_count=$(find . -type f -name "*.txt" | wc -l)
  echo "*.caption files: $caption_count"
  echo "*.txt files: $txt_count"
}

# Counts *.caption and *.txt files in each subdirectory individually.
count_captions_per_folder() {
  for dir in */; do
    echo "Directory: $dir"
    echo -n "*.caption files: "
    find "$dir" -type f -name "*.caption" | wc -l
    echo -n "*.txt files: "
    find "$dir" -type f -name "*.txt" | wc -l
  done
}

# open-webui
oui() {
  conda activate openwebui
  open-webui serve --port 6969 
}

llama() {
  ~/models/Meta-Llama-3-8B-Instruct.Q5_K_M.llamafile -cb -np 4 -a llama-3-8b --embedding --port 11434
}

alias gcs='git clone --recurse-submodules'

# Function to copy matching .caption files
copy_matching_caption_files() {
  # Define the target directory
  TARGET_DIR="$1"

  # Loop through each image file in the current directory
  for image_file in *.(jpg|jpeg|png|gif|bmp|tiff|webp|jxl); do
    # Check if the file exists (to handle cases where no files match the pattern)
    if [[ -f "$image_file" ]]; then
      # Extract the base name (without extension)
      base_name="${image_file%.*}"
      
      # Define the corresponding .caption file in the target directory
      caption_file="$TARGET_DIR/$base_name.caption"
      
      # Check if the .caption file exists
      if [[ -f "$caption_file" ]]; then
        # Copy the .caption file to the current directory
        cp "$caption_file" .
        echo "Copied $caption_file to the current directory."
      else
        echo "No matching .caption file for $image_file."
      fi
    fi
  done
}


# This script performs a text replacement operation in all .txt files within a specified directory.
# It takes three arguments:
# 1. target_dir: The directory containing the .txt files where the text replacement will occur.
# 2. search_text: The text string that needs to be replaced.
# 3. replace_text: The text string that will replace the search_text.
#
# The script uses a for loop to iterate through all .txt files in the target directory.
# It utilizes the 'sed' command to perform an in-place replacement of the search_text with the replace_text in each file.
# After processing all files, it prints a message indicating the completion of the text replacement operation.
replace_text_in_files() {
  local target_dir=$1
  local search_text=$2
  local replace_text=$3

  # Loop through all .txt files in the target directory
  for file in "$target_dir"/*.txt; do
    # Use sed to replace the text
    sed -i "s/$search_text/$replace_text/g" "$file"
  done

  echo "Text replacement complete in $target_dir!"
}

# Example usage:
# replace_text_in_files "/path/to/directory" "squishy (artist)" "by squishy (artist)"


# This script adds a specified prefix to the beginning of each text file in a given directory.
# Usage: inject_to_captions <directory> <prefix>
# Arguments:
#   <directory> - The directory containing the text files to be modified.
#   <prefix> - The prefix to be added to the beginning of each text file.
# The script checks if the specified directory exists and iterates over each text file in the directory.
# For each text file, it creates a temporary file with the modified content and then replaces the original file with the temporary file.
# If the directory does not exist, it prints an error message.
inject_to_captions() {
    local dir="$1"
    local prefix="$2"
    if [[ -d "$dir" ]]; then
        for file in "$dir"/*.txt; do
            if [[ -f "$file" ]]; then
                if ! grep -q "$prefix" "$file"; then
                    # Use a temporary file to store the modified content
                    local temp_file=$(mktemp)
                    echo "${prefix}, $(cat "$file")" > "$temp_file"
                    mv "$temp_file" "$file"
                    echo "Added '${prefix}, ' to the front of $file"
                else
                    echo "The tag '${prefix}' already exists in $file"
                fi
            fi
        done
    else
        echo "Directory $dir does not exist."
    fi
}

# Function to update git repositories in subdirectories
update_dir() {
    local target_dir="${1:-.}"
    
    # Check if there are any subdirectories
    if [[ -n "$(find "$target_dir" -mindepth 1 -maxdepth 1 -type d)" ]]; then
        for dir in "$target_dir"/*/; do
            if [[ -d "$dir" ]]; then
                (
                    cd "$dir" || return
                    # If the directory is a git repository, pull the latest changes
                    if [[ -d ".git" ]]; then
                        echo "Updating $(pwd)"
                        git pull
                    fi
                )
            fi
        done
    fi
}

export TOKENIZERS_PARALLELISM=false

alias grabber="Grabber-cli"

#export force_color_prompt=yes

chop_lora() {
    local input_file="$1"
    local base_name="${input_file:r}"  # Remove extension
    
    # Define presets and their corresponding vector strings
    declare -A presets=(
        ["ringdingding"] = "1,0,0,0,0,0,0,0,1,0,0,1,1,1,1,0,0,0,0,0,0"
        ["squeaker"]     = "1,0,0,0,0,0,0,0,1,0,0,0,1,1,1,0,0,0,0,0,0"
        ["heavylifter"]  = "1,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0"
        ["style1"]       = "1,0,0,0,1,1,0,1,1,0,0,0,1,1,1,1,1,1,0,0,0"
        ["style2"]       = "1,0,0,0,0,0,0,1,1,0,0,0,1,1,1,1,1,1,0,0,0"
        ["beeg"]         = "1,0,0,0,1,1,0,1,1,0,1,0,1,1,1,0,0,0,0,0,0"
        ["all"]          = "1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1"
        ["allin"]        = "1,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0"
        ["allmid"]       = "1,0,0,0,0,0,0,0,0,1,1,1,0,0,0,0,0,0,0,0,0"
        ["allout"]       = "1,0,0,0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1,1"
    )
    
    for preset in ${(k)presets}; do
        local output_file="${base_name}-${preset}.safetensors"
        local vector_string="${presets[$preset]}"
        echo "Generating $output_file"
        python ~/source/repos/resize_lora/chop_blocks.py "$input_file" "$vector_string" -o "$output_file"
    done
}

function swch() {
    if [ -z "$1" ]; then
        echo "Please provide a branch name."
        return 1
    fi
    branchname=$1
    git clean -fxd && git pull && git checkout $branchname
}

export COMFYUI_PATH="$HOME/ComfyUI"
export ZSH="$HOME/.oh-my-zsh"

ZSH_THEME="kade"
# CASE_SENSITIVE="true"
# HYPHEN_INSENSITIVE="true"
# DISABLE_MAGIC_FUNCTIONS="true"
# DISABLE_LS_COLORS="true"
# DISABLE_AUTO_TITLE="true"
# ENABLE_CORRECTION="true"
# COMPLETION_WAITING_DOTS="true"
# DISABLE_UNTRACKED_FILES_DIRTY="true"

plugins=(git autojump conda-env)

extract_iframes() {
    # Assign input arguments
    input_file="$1"
    scene_change_fraction="${2:-0.1}"

    # Get the base filename without extension
    base_name=$(basename "$input_file" .webm)

    # Run ffmpeg command
    /usr/bin/ffmpeg -i "$input_file" -f image2 -vf "select=eq(pict_type\,PICT_TYPE_I)*gt(scene\,$scene_change_fraction),showinfo" -fps_mode vfr "${base_name}-%06d.png"
}

convert_to_jxl() {
    local target_directory="$1"

    # Ensure the target directory exists
    if [[ ! -d "$target_directory" ]]; then
        echo "The specified directory does not exist: $target_directory" >&2
        return 1
    fi

    # Find all JPG, JPEG, and PNG files in the target directory and all subdirectories
    find "$target_directory" \( -name "*.jpg" -o -name "*.jpeg" -o -name "*.png" \) -type f | while read -r file; do
        input_path="$file"
        output_path="${file%.*}.jxl"

        # Convert to JXL using ImageMagick
        if magick convert "$input_path" "$output_path"; then
            echo "Converted: $input_path -> $output_path"
        else
            echo "Failed to convert $input_path" >&2
        fi
    done

    echo "Conversion complete."
}


convert_pxl_to_png() {
    local target_directory="$1"

    # Ensure the target directory exists
    if [[ ! -d "$target_directory" ]]; then
        echo "The specified directory does not exist: $target_directory" >&2
        return 1
    fi

    # Find all PXL files in the target directory and all subdirectories
    find "$target_directory" -type f -name "*.pxl" | while read -r file; do
        input_path="$file"
        output_path="${file%.pxl}.png"

        # Convert PXL to PNG using ImageMagick
        if magick convert "$input_path" "$output_path"; then
            echo "Converted: $input_path -> $output_path"
        else
            echo "Failed to convert $input_path" >&2
        fi
    done

    echo "Conversion complete."
}


seed() {
    local filePath="$1"
    python3 -c "
import safetensors, json
filePath = '$filePath'
print(json.loads(safetensors.safe_open(filePath, 'np').metadata().get('ss_seed', 'Not found')))
"
}

png2mp4() {
  ffmpeg -framerate 8 -pattern_type glob -i '*.png' -vf scale=512x512 -crf 28 \
  -c:v libx264 -pix_fmt yuv420p out.mp4
}


source $ZSH/oh-my-zsh.sh

export PATH=$PATH:$HOME/.local/bin:$HOME/source/repos/dataset-tools/target/x86_64-unknown-linux-gnu/release:$HOME/.cargo/bin:$HOME/miniconda3/bin:$HOME/toolkit:$HOME/db/redis-stable/src:$HOME/db/postgresql/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib
export COMFYUI_MODEL_PATH=/home/kade/ComfyUI/models

c_old() {
    cd ~/ComfyUI &&
    python3.12 main.py --listen 0.0.0.0 --preview-method taesd --use-pytorch-cross-attention --disable-xformers --fast
}

c() {
    cd ~/ComfyUI &&
    conda activate comfyui
    python main.py --listen 0.0.0.0 --preview-method taesd --use-pytorch-cross-attention --disable-xformers --front-end-version Comfy-Org/ComfyUI_frontend@latest --fast
}

alias t="tensorboard --logdir=$HOME/output_dir/logs"
alias nvim="vim"
alias rt="vim ~/.tmux.conf && echo \"Reloading tmux config\" && tmux source ~/.tmux.conf"
alias zr="vim ~/.zshrc && echo \"Reloading zsh config\" && source ~/.zshrc"
alias ta="tmux att"
alias ga="git add . && git commit -avs && git push"
alias gs="git status"
alias wd="git diff --word-diff-regex='[^,]+' --patience"

source /home/kade/.config/broot/launcher/bash/br

[ -f ~/.fzf.zsh ] && source ~/.fzf.zsh

alias ls='ls --color=always'

# >>> conda initialize >>>
# !! Contents within this block are managed by 'conda init' !!
__conda_setup="$('/home/kade/miniconda3/bin/conda' 'shell.zsh' 'hook' 2> /dev/null)"
if [ $? -eq 0 ]; then
    eval "$__conda_setup"
else
    if [ -f "/home/kade/miniconda3/etc/profile.d/conda.sh" ]; then
        . "/home/kade/miniconda3/etc/profile.d/conda.sh"
    else
        export PATH="/home/kade/miniconda3/bin:$PATH"
    fi
fi
unset __conda_setup
# <<< conda initialize <<<

unset CONDA_CHANGEPS1

function conda_prompt_info() {
  if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
    echo "(${CONDA_DEFAULT_ENV})"
  fi
}

display_custom_help