/*-
 * #%L
 * Fiji distribution of ImageJ for the life sciences.
 * %%
 * Copyright (C) 2015 - 2022 Fiji developers.
 * %%
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public
 * License along with this program.  If not, see
 * <http://www.gnu.org/licenses/gpl-3.0.html>.
 * #L%
 */
package fr.pasteur.iah.ccct;

import static fr.pasteur.iah.ccct.trackmate.CellContactDetectorFactory.KEY_CHANNEL_1;
import static fr.pasteur.iah.ccct.trackmate.CellContactDetectorFactory.KEY_CHANNEL_2;
import static fr.pasteur.iah.ccct.trackmate.CellContactDetectorFactory.KEY_CONTACT_SENSITIVITY;
import static fr.pasteur.iah.ccct.trackmate.CellContactDetectorFactory.KEY_SIGMA_FILTER;
import static fr.pasteur.iah.ccct.trackmate.CellContactDetectorFactory.KEY_THRESHOLD_1;

import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import javax.swing.ImageIcon;

import org.scijava.util.VersionUtils;

import fiji.plugin.trackmate.gui.GuiUtils;
import fiji.plugin.trackmate.util.TMUtils;
import fr.pasteur.iah.ccct.util.Thresholder;
import ij.CompositeImage;
import ij.IJ;
import ij.ImagePlus;
import ij.WindowManager;
import ij.plugin.PlugIn;
import net.imagej.ImgPlus;
import net.imagej.axis.Axes;
import net.imglib2.IterableInterval;
import net.imglib2.RandomAccessibleInterval;
import net.imglib2.algorithm.MultiThreaded;
import net.imglib2.algorithm.labeling.ConnectedComponentAnalysis;
import net.imglib2.algorithm.neighborhood.CenteredRectangleShape;
import net.imglib2.img.Img;
import net.imglib2.img.display.imagej.ImageJFunctions;
import net.imglib2.img.planar.PlanarImgFactory;
import net.imglib2.multithreading.SimpleMultiThreading;
import net.imglib2.type.NativeType;
import net.imglib2.type.logic.BitType;
import net.imglib2.type.numeric.RealType;
import net.imglib2.type.numeric.integer.UnsignedShortType;
import net.imglib2.util.Util;
import net.imglib2.view.IntervalView;
import net.imglib2.view.Views;

@SuppressWarnings( "deprecation" )
public class CCCT< T extends RealType< T > & NativeType< T > > implements PlugIn, CCCTProcessor, MultiThreaded
{

	public static final String PLUGIN_VERSION = VersionUtils.getVersion( CCCT.class );

	public static final String PLUGIN_NAME = "CCCT";

	public static final ImageIcon ICON = new ImageIcon( CCCT.class.getResource( "LogoCellContactAnalyzer-128x128.png" ) );;

	public static final ImageIcon ICON_SMALL = new ImageIcon( CCCT.class.getResource( "LogoCellContactAnalyzer-64x64.png" ) );;

	ImagePlus imp;

	private CCCTGui gui;

	private int numThreads;

	public CCCT()
	{
		setNumThreads( 10 );
	}

	@Override
	public void run( final String arg )
	{
		if ( null != arg && arg.length() > 0 )
		{
			imp = new CompositeImage( new ImagePlus( arg ) );
		}
		else
		{
			imp = WindowManager.getCurrentImage();
		}
		if ( null == imp )
		{
			IJ.error( PLUGIN_NAME + " v" + PLUGIN_VERSION, "Please open a multi-channel image first." );
			return;
		}
		if ( !imp.isVisible() )
		{
			imp.show();
		}

		gui = new CCCTGui( imp, this );
		GuiUtils.positionWindow( gui, imp.getWindow() );
		gui.setVisible( true );
	}

	@Override
	public void process( final boolean showContactImage, final boolean contactMask, final boolean contactLabels, final boolean trackLabels )
	{
		@SuppressWarnings( "unchecked" )
		final ImgPlus< T > img = TMUtils.rawWraps( imp );

		final Map< String, Object > settings = gui.getSettings();

		// In ImgLib2, dimensions are 0-based.
		final int channel1 = ( Integer ) settings.get( KEY_CHANNEL_1 ) - 1;
		final int channel2 = ( Integer ) settings.get( KEY_CHANNEL_2 ) - 1;
		final int contactSensitivity = ( Integer ) settings.get( KEY_CONTACT_SENSITIVITY );
		final double sigma = ( Double ) settings.get( KEY_SIGMA_FILTER );
		final double thresholdC1 = ( Double ) settings.get( KEY_THRESHOLD_1 );
		final double thresholdC2 = ( Double ) settings.get( KEY_THRESHOLD_1 );

		final RandomAccessibleInterval< T > im1;
		final RandomAccessibleInterval< T > im2;

		final int cDim = img.dimensionIndex( Axes.CHANNEL );
		if ( cDim < 0 )
		{
			im1 = img;
			im2 = img;
		}
		else
		{
			im1 = Views.hyperSlice( img, cDim, channel1 );
			im2 = Views.hyperSlice( img, cDim, channel2 );
		}

		final PlanarImgFactory< T > factory = new PlanarImgFactory<>( Util.getTypeFromInterval( img ) );
		final Img< T > out = factory.create( im1 );

		int timeDim = img.dimensionIndex( Axes.TIME );
		if ( timeDim >= 0 )
			if ( cDim >= 0 && timeDim > cDim )
				timeDim--;

		final int td = timeDim;
		final int nFrames = imp.getNFrames();

		{
			gui.setProgressStatus( "Contact image" );
			final AtomicInteger ai = new AtomicInteger( 0 );
			final Thread[] threads = SimpleMultiThreading.newThreads( numThreads );
			for ( int i = 0; i < threads.length; i++ )
			{
				threads[ i ] = new Thread( "CCCT Thread Contact Image" )
				{
					@Override
					public void run()
					{
						for ( int frame = ai.getAndIncrement(); frame < nFrames; frame = ai.getAndIncrement() )
						{
							final ContactImgGenerator< T > algo = new ContactImgGenerator< T >(
									Views.hyperSlice( im1, td, frame ),
									Views.hyperSlice( im2, td, frame ),
									Views.hyperSlice( out, out.numDimensions() - 1, frame ),
									thresholdC1, thresholdC2, contactSensitivity, sigma );

							if ( !algo.checkInput() || !algo.process() )
							{
								System.err.println( algo.getErrorMessage() );
								return;
							}

							gui.setProgress( ( 1.0 + ai.get() ) / nFrames );
						}
					}
				};
			}
			SimpleMultiThreading.startAndJoin( threads );

			if ( showContactImage )
			{
				final ImagePlus contacts = ImageJFunctions.wrap( out, "Contacts" );
				contacts.setCalibration( imp.getCalibration().copy() );
				contacts.setDimensions( 1, imp.getNSlices(), imp.getNFrames() );
				contacts.show();
				contacts.resetDisplayRange();
				IJ.run( "Fire" );
			}
		}
		/*
		 * Generate binary mask.
		 */

		if ( contactMask || contactLabels )
		{
			final PlanarImgFactory< BitType > maskFactory = new PlanarImgFactory<>( new BitType() );
			final Img< BitType > mask = maskFactory.create( im1 );
			final T valTreshold = out.firstElement().createVariable();
			valTreshold.setZero();

			gui.setProgressStatus( "Contact masks" );
			final AtomicInteger ai = new AtomicInteger( 0 );
			final Thread[] threads = SimpleMultiThreading.newThreads( numThreads );
			for ( int i = 0; i < threads.length; i++ )
			{
				threads[ i ] = new Thread( "CCCT Thread Contact Mask" )
				{
					@Override
					public void run()
					{
						for ( int frame = ai.getAndIncrement(); frame < nFrames; frame = ai.getAndIncrement() )
						{
							final IterableInterval< BitType > maskSlice = Views.hyperSlice( mask, mask.numDimensions() - 1, frame );
							final IntervalView< T > slice = Views.hyperSlice( out, out.numDimensions() - 1, frame );
							Thresholder.threshold( slice, maskSlice, valTreshold, true, 1 );
							gui.setProgress( ( 1.0 + ai.get() ) / nFrames );
						}
					}
				};
			}
			SimpleMultiThreading.startAndJoin( threads );

			if ( contactMask )
			{
				final ImagePlus masks = ImageJFunctions.wrap( mask, "ContactMasks" );
				masks.setCalibration( imp.getCalibration() );
				masks.setDimensions( 1, imp.getNSlices(), imp.getNFrames() );
				masks.show();
			}

			if ( contactLabels )
			{
				gui.setProgressStatus( "Contact labels" );
				gui.setProgress( 0. );
				final Img< UnsignedShortType > lbl = Util.getArrayOrCellImgFactory( mask, new UnsignedShortType() ).create( mask );

				if ( trackLabels && imp.getNFrames() > 1 )
					/*
					 * Connect in all dimensions, including time. This will
					 * effectively track contacts, provided they overlap from
					 * one frame to another.
					 */
					ConnectedComponentAnalysis.connectedComponents( mask, lbl );
				else
				{
					/*
					 * Skip the time dimension when connecting mask.
					 */
					final int numDimensions = mask.numDimensions();
					final int[] span = new int[ numDimensions ];
					Arrays.fill( span, 1 );
					if ( imp.getNSlices() > 1 )
						span[ 3 ] = 0;
					else
						span[ 2 ] = 0;

					ConnectedComponentAnalysis.connectedComponents( mask, lbl, new CenteredRectangleShape( span, false ) );
				}

				// For display range.
				int maxLbl = -1;
				for ( final UnsignedShortType pixel : lbl )
					if ( pixel.get() > maxLbl )
						maxLbl = pixel.get();

				final ImagePlus labels = ImageJFunctions.wrap( lbl, "ContactLabels" );
				labels.setCalibration( imp.getCalibration().copy() );
				labels.setDimensions( 1, imp.getNSlices(), imp.getNFrames() );
				labels.setDisplayRange( 0, maxLbl );
				labels.show();
				IJ.run( "3-3-2 RGB" );
			}
		}
	}

	@Override
	public void setNumThreads()
	{
		this.numThreads = Runtime.getRuntime().availableProcessors();
	}

	@Override
	public void setNumThreads( final int numThreads )
	{
		this.numThreads = numThreads;
	}

	@Override
	public int getNumThreads()
	{
		return numThreads;
	}
}