#!/bin/bash

set -e

script="$0"

count=0
while [ -L "${script}" ]
do
	script=$(readlink ${script})
	count=$(expr ${count} + 1)
	if [ ${count} -gt 100 ]
	then
		echo "Too many symbolic links"
		exit 1
	fi
done
scriptdir=$(dirname ${script})
scriptname=$(basename ${script})

. "$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 ! 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 ! pdf_tools_installed; then
	cat <<EOF

${color_warning}NOTE: Pdf tools are missing; Will not generate one single result pdf.${color_reset}
EOF
else
	pdf_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 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/darktable-noiseprofile"
	if [ ! -x "$scriptdir/darktable-noiseprofile" ]; then
		echo "ERROR: Could not find darktable-noiseprofile tool in $scriptdir"
		exit 1
	fi

	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" --apply-custom-presets false --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 darktable-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"
	set xlabel "brightness"
	set ylabel "standard deviation"
	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"
	set xlabel "brightness"
	set ylabel "standard deviation"
	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",				\
	       1 w l lw 4 title "ideal"
	set title "Flat histogram - $title"
	set xlabel "brightness"
	set ylabel ""
	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"
	set xlabel ""
	set ylabel ""
	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 exposure.
    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 exposure.
    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 exposure.
    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

	printf "$maker\t$model\t$model iso $iso\t$iso\t$a0, $a1, $a2\t$b0, $b1, $b2\n" >> "$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"
}

txt2json(){
  local txt json end n_makers n_models n_profiles maker model profile comment

  txt="$1"
  json="$2"

(
  cat <<EOF
{
  "version": 0,
  "noiseprofiles": [
EOF

  n_makers=$(cut -f 1 "$txt" | sort -u | wc -l)
  cut -f 1 "$txt" | sort -u | while read maker
  do
    n_makers=$((n_makers-1))
    cat <<EOF
    {
      "maker": "$maker",
      "models": [
EOF
    n_models=$(awk -F "\t" "{ if(\$1 == \"$maker\"){ print \$2; } }" "$txt" | sort -u | wc -l)
    awk -F "\t" "{ if(\$1 == \"$maker\"){ print \$2; } }" "$txt" | sort -u | while read model
    do
      n_models=$((n_models-1))
      comment=$(echo $model | tr [:upper:] [:lower:])" contributed by "$(whoami)
      cat <<EOF
        {
          "comment": "$comment",
          "model": "$model",
          "profiles": [
EOF
      n_profiles=$(awk -F "\t" "{ if(\$1 == \"$maker\" && \$2 == \"$model\"){ print \$1; } }" "$txt" | wc -l)
      awk -F "\t" "{ if(\$1 == \"$maker\" && \$2 == \"$model\"){ printf(\"\\\"name\\\": \\\"%s\\\", \\\"iso\\\": %s, \\\"a\\\": [%s], \\\"b\\\": [%s]\n\", \$3, \$4, \$5, \$6); } }" "$txt" | sort -V | while read profile
      do
        n_profiles=$((n_profiles-1))
        end=","
        if [ $n_profiles -eq 0 ]; then
          end=""
        fi
        echo "            {$profile}$end"
      done
      end=","
      if [ $n_models -eq 0 ]; then
        end=""
      fi
      cat <<EOF
          ]
        }$end
EOF
    done
    end=","
    if [ $n_makers -eq 0 ]; then
      end=""
    fi
    cat <<EOF
      ]
    }$end
EOF
  done

  cat <<EOF
  ]
}
EOF
) > "$json"

}

# --------------------------------------------------------------------
# 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_exposure "$image" "$jpeg_export"; then
				if [ "$incorrect_isos" ]; then
					incorrect_isos="$incorrect_isos $iso"
				else
					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 "    - $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  remove previous presets

echo
echo "===> Prepare profiling job"

presets_txt="$profiling_dir/presets.txt"
presets_json="$profiling_dir/presets.json"

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

# 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_txt" "$files_list"
	done
done

# turn presets.txt into presets.json
txt2json "$presets_txt" "$presets_json"
rm -f "$presets_txt"

#add all pdfs into one
if [ "$pdf_enabled" ]; then
	resultpdf="$profiling_dir"/noise_result.pdf
	echo
	echo "===> Prepare final pdf"

	for iso in $iso_settings; do
	images=$(get_var "images_$iso")
		for image in $images; do
			pdf="${image%.*}.pdf"
			pdfflat="${image%.*}_flat.pdf"
			flist="$flist $pdf $pdfflat"
		done
	done
	echo Filelist: $flist
	pdfcat $resultpdf $flist
fi

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

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

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

To test the noise profiles locally, run:

EOF

dt_version=$(get_darktable_version)
if cmp_darktable_version "$dt_version" "-lt" "1.6.3"; then
	cat <<EOF
  /path/to/darktable-1.6.3-or-later/darktable --noiseprofiles $presets_json

${color_warning}CAUTION: The "darktable" executable available from your system's \$PATH
has version $dt_version. However, version 1.6.3 or later is required to
test the profile. The command line shown above must be changed to point
to your darktable 1.6.3+ binary.${color_reset}
EOF
else
	cat <<EOF
  darktable --noiseprofiles $presets_json
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
