livia-test/Livia/ViewModels/ImageRotationViewerControlViewModel.cs

using System.Collections.ObjectModel;
using System.IO;
using System.Windows;
using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Messaging;
using Livia.Models;
using Livia.Models.Data;
using Livia.Utility;
using Livia.Utility.DependencyInjection;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;

namespace Livia.ViewModels;

public enum ImageRotationViewerDataType
{
    Dicom2D,
    Dicom3D,
    PlainImage
}

public class ImageRotationViewerControlViewModel : ObservableRecipient, IRecipient<ImageRotationViewerIndexChangeMessage>, IRecipient<RoiTabChangedMessage>, ILiviaModuleViewModel, IRecipient<MaskChangedMessage>
{
    public IImageSeries StructureImageSeries { get; private set; }
    public IImageSeries AtlasMaskImageSeries { get; private set; }
    public ObservableCollection<RoiExpanderControlViewModel> CurrentRoiCollection { get; } = [];
    public int MaxIndex { get => _maxIndex; private set => SetProperty(ref _maxIndex, value); }
    public bool ShowHelperTip { get => _showHelperTip;  set => SetProperty(ref _showHelperTip, value); }

    public int CurrentIndex
    {
        get => _currentIndex;
        private set
        {
            SetProperty(ref _currentIndex, value);
            StructureImageSeries.SetImageIndex(CurrentIndex - 1);
            AtlasMaskImageSeries.SetImageIndex(CurrentIndex - 1);
            foreach (RoiExpanderControlViewModel viewModel in _roiDictionary.SelectMany(keyValuePair => keyValuePair.Value))
            {
                viewModel.MaskImageSeries.SetImageIndex(CurrentIndex - 1);
            }
        }
    }
    public string DisplayName { get => _displayName; set => SetProperty(ref _displayName, value); }
    public string AtlasMaskLoadPath { get; set; } = string.Empty;
    public string StructLoadPath { get; set; } = string.Empty;
    public string ImageIndexSyncKey { get; set; } = string.Empty;
    public ImageRotationViewerDataType StructDataType { get; set; }
    public ColorBarControlViewModel? ColorBarControlViewModel { get => _colorBarControlViewModel; set => SetProperty(ref _colorBarControlViewModel, value); }
    public string AtlasId { get; set; } = string.Empty;
    //TEMP
    public static string LoadAtlasName = "ASL-cCBF";
    public double WindowCenterModifier
    {
        get => StructureImageSeries.WindowCenterModifier;
        set
        {
            StructureImageSeries.WindowCenterModifier = value;
            //refresh image
            WeakReferenceMessenger.Default.Send(new ImageRotationViewerIndexChangeMessage(ImageIndexSyncKey, CurrentIndex));
        }
    }
    public double WindowWidthModifier
    {
        get => StructureImageSeries.WindowWidthModifier;
        set
        {
            StructureImageSeries.WindowWidthModifier = value;
            //refresh image
            WeakReferenceMessenger.Default.Send(new ImageRotationViewerIndexChangeMessage(ImageIndexSyncKey, CurrentIndex));
        }
    }

    private ColorBarControlViewModel? _colorBarControlViewModel;
    private int _maxIndex;
    private bool _showHelperTip;
    private int _currentIndex;
    private string _displayName = string.Empty;

    private readonly Dictionary<int, IEnumerable<RoiExpanderControlViewModel>> _roiDictionary = new();
    private readonly ILogger _logger;
    private readonly IDataBlockLoader _dataBlockLoader;
    private DataBlock? _currentDataBlock;

    public ImageRotationViewerControlViewModel(ILogger logger, IDataBlockLoader dataBlockLoader)
    {
        _logger = logger;
        _dataBlockLoader = dataBlockLoader;
        WeakReferenceMessenger.Default.RegisterAll(this);

        //init so compiler won't complain
        StructureImageSeries = new PlainImageSeries();
        AtlasMaskImageSeries = new PlainImageSeries();
    }

    public void NextImage()
    {
        if (CurrentIndex >= MaxIndex)
            return;

        _logger.LogInformation("Next image, showing {i}", CurrentIndex);
        WeakReferenceMessenger.Default.Send(new ImageRotationViewerIndexChangeMessage(ImageIndexSyncKey, CurrentIndex + 1));
    }

    public void PrevImage()
    {
        if (CurrentIndex <= 1)
            return;

        _logger.LogInformation("Prev image, showing {i}", CurrentIndex - 2);
        WeakReferenceMessenger.Default.Send(new ImageRotationViewerIndexChangeMessage(ImageIndexSyncKey, CurrentIndex - 1));
    }

    public async Task LoadData(DataBlock dataBlock)
    {
        _currentDataBlock = dataBlock;
        if (!string.IsNullOrEmpty(AtlasId))
        {
            AdditionalInfoJson additionalInfo = _dataBlockLoader.GetSharedData<AdditionalInfoJson>("additionalInfo");
            int findIndex = additionalInfo.AtlasInfoListDict[LoadAtlasName].FindIndex(item => item.AtlasName == AtlasId);
            if (findIndex < 0)
            {
                _logger.LogWarning("Unknown atlas id: {id}", AtlasId);
            }
            else
            {
                ColorBarControlViewModel = ActivatorUtilities.CreateInstance<ColorBarControlViewModel>(ServiceProviderFactory.ServiceProvider);
                AtlasInfo info = additionalInfo.AtlasInfoListDict[LoadAtlasName][findIndex];
                ColorBarControlViewModel.ColorBarRangeMin = info.Range[0];
                ColorBarControlViewModel.ColorBarRangeMax = info.Range[1];
                ColorBarControlViewModel.ColorBarUnit = info.Unit;
            }
        }

        string loadPath = AtlasMaskLoadPath.Replace("mask", $"{LoadAtlasName}_mask");
        if (!string.IsNullOrEmpty(AtlasMaskLoadPath))
        {
            await AtlasMaskImageSeries.LoadData(Path.Combine(dataBlock.ResultPath, loadPath));
        }

        if (!string.IsNullOrEmpty(StructLoadPath))
        {
            await StructureImageSeries.LoadData(Path.Combine(dataBlock.ResultPath, StructLoadPath));
        }

        //_roiDictionary.Clear();
        
        // set to middle
        MaxIndex = StructureImageSeries.Count;
        WeakReferenceMessenger.Default.Send(new ImageRotationViewerIndexChangeMessage(ImageIndexSyncKey, MaxIndex / 2));
    }

    public void Init()
    {
        //we want all image series that share the same load path share one IImage series. So when we change WindowWidth, everything is synced.
        string structureImageSeriesKey = $"structureImageSeries{StructLoadPath}";
        if (_dataBlockLoader.SharedDataContainsKey(structureImageSeriesKey))
        {
            StructureImageSeries = _dataBlockLoader.GetSharedData<IImageSeries>(structureImageSeriesKey);
        }
        else
        {
            StructureImageSeries = BuildImageSeries(StructDataType);
            _dataBlockLoader.SaveSharedData(structureImageSeriesKey, StructureImageSeries);
        }

        AtlasMaskImageSeries = BuildImageSeries(ImageRotationViewerDataType.PlainImage);
    }

    private static IImageSeries BuildImageSeries(ImageRotationViewerDataType type)
    {
        return type switch
        {
            ImageRotationViewerDataType.Dicom2D => new Dicom2DImageSeries(),
            ImageRotationViewerDataType.Dicom3D => new Dicom3DImageSeries(),
            ImageRotationViewerDataType.PlainImage => new PlainImageSeries(),
            _ => throw new ArgumentOutOfRangeException(nameof(type), type, null)
        };
    }

    public void SetDefaultWindowLevel()
    {
        StructureImageSeries.WindowWidthModifier = 1;
        StructureImageSeries.WindowCenterModifier = 1;

        //refresh image
        WeakReferenceMessenger.Default.Send(new ImageRotationViewerIndexChangeMessage(ImageIndexSyncKey, CurrentIndex));
    }

    public void AddRoiListToDictionary(int tabIndex, IEnumerable<RoiExpanderControlViewModel> list)
    {
        _roiDictionary.TryAdd(tabIndex, list);
    }

    public void Receive(ImageRotationViewerIndexChangeMessage message)
    {
        if (ImageIndexSyncKey == string.Empty || ImageIndexSyncKey != message.Value.Item1)
            return;

        CurrentIndex = message.Value.Item2;
    }

    public void Receive(RoiTabChangedMessage message)
    {
        CurrentRoiCollection.Clear();
        if (_roiDictionary.TryGetValue(message.Value, out IEnumerable<RoiExpanderControlViewModel>? list))
        {
            CurrentRoiCollection.AddRange(list.OrderBy(item => item.ZIndex));
            DisplayName = (string)(Application.Current.TryFindResource($"RoiTabHeader{message.Value}") ?? "");
        }
        else if (CurrentRoiCollection.Count == 0 && _roiDictionary.Count > 0)
        {
            // we are not displaying anything, so just pick something to show
            KeyValuePair<int, IEnumerable<RoiExpanderControlViewModel>> pair = _roiDictionary.First();
            CurrentRoiCollection.AddRange(pair.Value.OrderBy(item => item.ZIndex));
            DisplayName = (string)(Application.Current.TryFindResource($"RoiTabHeader{pair.Key}") ?? "");
        }
        else
        {
            //nothing to show
            return;
        }

        //refresh
        CurrentIndex = _currentIndex;
    }

    public void ClearData()
    {
        StructureImageSeries.Dispose();
        AtlasMaskImageSeries.Dispose();
        _roiDictionary.Clear();
        CurrentRoiCollection.Clear();
        WeakReferenceMessenger.Default.UnregisterAll(this);
    }

    public void Receive(MaskChangedMessage message)
    {
        if (_currentDataBlock == null)
            return;

        Application.Current.Dispatcher.Invoke(() => LoadData(_currentDataBlock));
    }
}