#!/bin/sh

set -e

scriptname=$(basename $0)
scriptdir=$(cd $(dirname $0) && pwd)

. "$scriptdir/subr.sh"

# --------------------------------------------------------------------
# Read command line arguments.
# --------------------------------------------------------------------

usage() {
	cat <<EOF
$scriptname [OPTION]...

Available options:
    -C         Skip input images correctness checks.
    -d DIR     Use DIR to read and write profiling input RAW images.
    -h         Display this usage.
    -i ISO     Limit capture to ISO sensitivity; can be specified multiple
               times for multiple ISO settings.
    -K         Keep PFM exported files (caution: takes a lot of disc space).
    -P         Re-do tethered input images capture.
    -p SEC     Wait SEC seconds between each shot; useful if using a flash
               for instance.
EOF
}

while getopts ":Cd:hi:KPp:" opt; do
	case $opt in
	C)
		skip_correctness_check=1
		;;
	d)
		profiling_dir=$OPTARG
		;;
	h)
		usage
		exit 0
		;;
	i)
		iso_settings=$(add_to_list "$iso_settings" $OPTARG)
		;;
	K)
		keep_pfms=1
		;;
	P)
		force_profiling_shots=1
		;;
	p)
		pause_between_shots=$OPTARG
		;;
	esac
done
shift $((OPTIND-1))

# Sort user-specified ISO settings.
if [ "$iso_settings" ]; then
	iso_settings=$(sort_iso_list $iso_settings)
fi

# Check for required tools.
echo "===> Check for required tools"
missing_tool=0

if ! image_info_tools_installed; then
	missing_tool=1
fi
if ! image_export_tools_installed; then
	missing_tool=1
fi
if ! profiling_tools_installed; then
	missing_tool=1
fi
if ! database_tools_installed; then
	missing_tool=1
fi
if ! internal_tools_available; then
	missing_tool=1
fi
if ! tethering_tools_installed; then
	cat <<EOF

${color_warning}NOTE: Tethering tools are missing; you'll need to provide input images
yourself.${color_reset}
EOF
else
	tethering_enabled=1
fi

if [ "$missing_tool" = "1" ]; then
	exit 1
fi

# --------------------------------------------------------------------
# Internal functions.
# --------------------------------------------------------------------

# CAUTION: This function uses the following global variables:
#     o  scriptdir
#     o  keep_pfms

profile_image() {
	local image preset files_list maker model iso xmp		\
	 noiseprofile floatdump pfm dat title fit pdf			\
	 flat_dat curves_dat flat_pdf a0 a1 a2 b0 b1 b2

	tool_installed darktable-cli
	tool_installed gnuplot

	image="$1"
	presets="$2"
	files_list="$3"

	maker=$(get_image_camera_maker "$image")
	model=$(get_image_camera_model "$image")
	iso=$(get_image_iso "$image")

	xmp="$scriptdir/profiling-shot.xmp"
	noiseprofile="$scriptdir/noiseprofile"
	floatdump="$scriptdir/floatdump"

	echo
	echo "===> Profile image for \"$maker - $model - $iso ISO\""

	pfm="${image%.*}.pfm"
	if [ ! -f "$pfm" -o "$image" -nt "$pfm" -o "$xmp" -nt "$pfm" -o "$0" -nt "$pfm" ]; then
		echo "--> Converting $image (ISO $iso)"
		rm -f "$pfm"
		darktable-cli "$image" "$xmp" "$pfm" --core --conf plugins/lighttable/export/iccprofile=image --conf plugins/lighttable/export/style=none
	else
		echo "--> Skip $image (ISO $iso); output up-to-date"
	fi

	echo "--> Run noiseprofile"
	dat="${pfm%.pfm}.dat"
	$noiseprofile "$pfm" > "$dat"

	echo "--> Plotting $pfm"
	title="$maker, $model, $iso ISO ($(basename "${pfm%.pfm}"))"
	fit="${pfm%.pfm}.fit"
	pdf="${pfm%.pfm}.pdf"
	gnuplot 2>/dev/null <<EOF
	set term pdf
	set print "$fit"
	set output "$pdf"
	set fit logfile "/dev/null"
	set title "Histogram - $title"
	plot "$dat" u 1:(log(\$5)) w l lw 4 title "red",		\
	  '' u 1:(log(\$6)) w l lw 4 title "green",			\
	  '' u 1:(log(\$7)) w l lw 4 title "blue"

	min(x,y) = (x < y) ? x : y
	max(x,y) = (x > y) ? x : y
	f1(x) = a1*x + b1
	f2(x) = a2*x + b2
	f3(x) = a3*x + b3
	a1=0.1;b1=0.01;
	a2=0.1;b2=0.01;
	a3=0.1;b3=0.01;
	set xrange [0:0.35]
	fit f1(x) "$dat" u 1:(\$2**2):(1/max(0.01, \$5)) via a1,b1
	set xrange [0:0.8]
	fit f2(x) "$dat" u 1:(\$3**2):(1/max(0.01, \$6)) via a2,b2
	set xrange [0:0.5]
	fit f3(x) "$dat" u 1:(\$4**2):(1/max(0.01, \$7)) via a3,b3

	set xrange [0:1]
	set title "Noise levels - $title"
	plot "$dat" u 1:2 w l lw 4 title "red",				\
	  '' u 1:3 w l lw 4 title "green",				\
	  '' u 1:4 w l lw 4 title" blue",				\
	  '' u 1:(sqrt(f1(\$1))) w l lw 2 lt 1 title "red (fit)",	\
	  '' u 1:(sqrt(f2(\$1))) w l lw 2 lt 2 title "green (fit)",	\
	  '' u 1:(sqrt(f3(\$1))) w l lw 2 lt 3 title "blue (fit)"

	print a1, a2, a3, b1, b2, b3
EOF

	# Fitted parametric curves:
	echo "--> Fitted parametric curves"
	flat_dat="${pfm%.pfm}_flat.dat"
	curves_dat="${pfm%.pfm}_curves.dat"
	$noiseprofile $pfm -c $(cat $fit) > $flat_dat 2> $curves_dat

	# Data based histogram inversion:
	# $noiseprofile $pfm -h $dat > $flat_dat 2> $curves_dat
	echo "--> Flattened $pfm"
	flat_pdf="${pfm%.pfm}_flat.pdf"
	gnuplot 2>/dev/null << EOF
	set term pdf
	set output "$flat_pdf"
	set title "Flat noise levels - $title"
	plot "$flat_dat" u 1:2 w l lw 4 title "red",			\
	  '' u 1:3 w l lw 4 title "green",				\
	  '' u 1:4 w l lw 4 title "blue"
	set title "Flat histogram - $title"
	plot "$dat" u 1:(log(\$5)) w l lw 4 title "red",		\
	  '' u 1:(log(\$6)) w l lw 4 title "green",			\
	  '' u 1:(log(\$7)) w l lw 4 title "blue"
	set title "Conversion curves - $title"
	plot "$curves_dat" u 0:1 w l lw 4 title "red",			\
	  '' u 0:2 w l lw 4 title "green",				\
	  '' u 0:3 w l lw 4 title "blue"
EOF
	# Output preset for dt:
	echo "--> Save generated preset"
	a0=$(cat $fit | cut -f1 -d' ')
	a1=$(cat $fit | cut -f2 -d' ')
	a2=$(cat $fit | cut -f3 -d' ')
	b0=$(cat $fit | cut -f4 -d' ')
	b1=$(cat $fit | cut -f5 -d' ')
	b2=$(cat $fit | cut -f6 -d' ')

	case "$a1" in
	-*)
		cat <<EOF

${color_error}ERROR: Incorrect green channel.

Possible reason:
    o  The input RAW image is bad regarding lighting or exposition.
    o  You may have set a default output profile in darktable which is
       unsuitable for noise profiling.

If the latter reason applies to you, change it back to "image settings"
and re-run this script.${color_reset}
EOF
		rm -f $dat $flat_dat $curves_dat $fit "$pfm"
		return 1
		;;
	esac

	case "$a0" in
	-*)
		cat <<EOF

${color_warning}WARNING: Incorrect red channel.

Possible reason:
    o  The input RAW image is bad regarding lighting or exposition.
    o  You may have set a default output profile in darktable which is
       unsuitable for noise profiling.

If the latter reason applies to you, change it back to "image settings"
and re-run this script.

Meanwhile, this channel will be ignored.${color_reset}
EOF
		a0="0.0"
		b0="0.0"
		;;
	esac

	case "$a2" in
	-*)
		cat <<EOF

${color_warning}WARNING: Incorrect blue channel.

Possible reason:
    o  The input RAW image is bad regarding lighting or exposition.
    o  You may have set a default output profile in darktable which is
       unsuitable for noise profiling.

If the latter reason applies to you, change it back to "image settings"
and re-run this script.

Meanwhile, this channel will be ignored.${color_reset}
EOF
		a2="0.0"
		b2="0.0"
		;;
	esac

	echo "{\"$model iso $iso\",       \"$maker\",      \"$model\",              $iso,         {$a0, $a1, $a2}, {$b0, $b1, $b2}}," >> "$presets"

	# Clean unused files.
	rm -f $dat $flat_dat $curves_dat $fit
	if [ "$keep_pfms" != "1" ]; then
		rm -f "$pfm"
	fi

	echo "$(basename "$pdf")" >> "$files_list"
	echo "$(basename "$flat_pdf")" >> "$files_list"
}

# --------------------------------------------------------------------
# Part #1: Get profiling shots.
# --------------------------------------------------------------------

# If the user didn't specified a profiling shots directory, use a
# default one.
#
# Defaults to /var/tmp/darktable-noise-profiling/$camera/profiling.

auto_set_profiling_dir "-d"

# Check the profiling shots to see if there's at least one shot for
# each ISO settings, either specified by the user, or supported by the
# camera.
list_input_images

# Take the required shots.
if [ "$tethering_enabled" ]; then
	auto_capture_images
elif [ -z "$iso_settings" ]; then
	iso_settings=$images_for_iso_settings
fi

files_list="$profiling_dir/output-files-list.txt"
rm -f "$files_list"

echo
if [ "$skip_correctness_check" = "1" ]; then
	echo "===> Summing up profiling RAW images + Jpeg export"
else
	echo "===> Checking profiling RAW images correctness + Jpeg export"
fi
for iso in $iso_settings; do
	echo "--> ISO $iso:"
	images=$(get_var "images_$iso")
	for image in $images; do
		echo "    $image"

		# Export RAW file to jpeg file, then make a thumbnail
		# from this jpeg. The large jpeg file is used to perform
		# correctness checks. The small jpeg file added to the
		# final tarball to have a preview of the original RAW
		# file and a copy of the Exif metadata.
		jpeg_export="${image%.*}-large.jpg"
		jpeg_thumb="${image%.*}-thumb.jpg"
		if [ ! -f "$jpeg_export" -o "$image" -nt "$jpeg_export" -o "$0" -nt "$jpeg_export" ]; then
			export_large_jpeg "$image" "$jpeg_export"
			export_thumbnail "$jpeg_export" "$jpeg_thumb"
		fi
		echo "$(basename "$jpeg_thumb")" >> $files_list

		if [ "$skip_correctness_check" != "1" ]; then
			if ! check_exposition "$image" "$jpeg_export"; then
				if [ "$incorrect_isos" ]; then
					incorrect_isos="$iso"
				else
					incorrect_isos="$incorrect_isos $iso"
				fi
			fi
		fi
	done
done

if [ "$incorrect_isos" ]; then
	cat <<EOF

The following ISO settings have incorrect images:
EOF
	for iso in $incorrect_isos; do
		echo "    o  $iso ISO"
	done
	cat <<EOF

Please read the error messages associated with each file above, check
your subject and re-do the shots.
EOF
	exit 1
fi

# --------------------------------------------------------------------
# Part #2: Profile all images.
# --------------------------------------------------------------------

# First, prepare the working directory:
#   o  build tools
#   o  copy darktable database
#   o  remove previous presets

echo
echo "===> Prepare profiling job"

presets="$profiling_dir/presets.txt"

echo "--> Copy darktable library for testing purpose"
database="$profiling_dir/library.db"
database_orig="$HOME/.config/darktable/library.db"
if [ ! -f "$database_orig" ]; then
	cat <<EOF
${color_error}ERROR: Please run darktable at least once to create an initial library
or copy some valid library to $database manually before running this
script. This is needed to setup a testing environment for your new
presets.${color_reset}
EOF
	exit 1
fi

cp "$database_orig" "$database"

echo "--> Remove previous presets"
rm -f "$profiling_dir"/*.pdf
rm -f "$profiling_dir"/*.fit
rm -f "$profiling_dir"/*.dat
rm -f "$presets"

# Now, for each shots, profile it.
cat <<EOF
--> Ready to profile images

NOTE: This process takes some time and a lot of memory and disc space
(up-to several gigabytes, depending on the number of ISO settings and
the size of the RAW files.
EOF

for iso in $iso_settings; do
	images=$(get_var "images_$iso")
	for from in $images; do
		profile_image "$from" "$presets" "$files_list"
	done
done

# Insert presets in the copied database.
echo
echo "===> Record presets in library for testing purpose"
"$scriptdir"/add-profile -l "$database" -P "[test] " "$presets"

# Prepare a tarball ready to be sent to the darktable team.
echo
echo "===> Prepare final tarball"
tarball="$profiling_dir"/dt-noiseprofile-$(date +'%Y%m%d').tar.gz
cd "$profiling_dir"
tar -cf - "$(basename "$presets")" $(cat "$(basename "$files_list")") | gzip > "$(basename "$tarball")"
cd -

echo "--> Cleanup"
rm -f "$files_list"

echo ""
echo "${color_ok}== Noise profiling done! ==${color_reset}"
cat <<EOF

Presets were inserted into $database. to test them locally, run:

EOF

dt_version=$(get_darktable_version)
if cmp_darktable_version "$dt_version" "-lt" "1.2"; then
	cat <<EOF
  /path/to/darktable-1.2-or-later/darktable --library $database

${color_warning}CAUTION: The "darktable" executable available from your system's \$PATH
has version $dt_version. However, version 1.2 or later is required to
test the profile. The command line shown above must be changed to point
to your darktable 1.2+ binary.${color_reset}
EOF
else
	cat <<EOF
  darktable --library $database
EOF
fi

cat <<EOF

If you're happy with the results, post the following file to us:

  $tarball

If not, probably something went wrong. It's a good idea to get in touch
so we can help you sort it out.
EOF
