Dataset.Transform.cs

using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Text.Json;
using JSONstat.IO;
using JSONstat.Transform;

namespace JSONstat.Models;

/// <summary>
/// Transforming methods on <see cref="Dataset"/>: the non-deprecated Toolkit
/// replacements <c>Dice</c> and <c>Transform</c>, plus <c>Unflatten</c>. The
/// deprecated <c>Slice</c> and <c>toTable</c> are intentionally absent.
/// </summary>
public sealed partial class Dataset
{
    /// <summary>
    /// Creates a subset by filtering dimension categories.
    /// </summary>
    public Dataset Dice(DiceFilter? filter = null, DiceOptions? options = null)
    {
        options ??= new DiceOptions();
        filter ??= DiceFilter.Empty();
        EnsureLinked();

        var ids = Id ?? Array.Empty<string>();
        var sizes = Sizes;
        int d = ids.Length;

        var keptIds = new string[d][];
        var keptOldPos = new int[d][];

        for (int i = 0; i < d; i++)
        {
            var dim = _orderedDimensions[i];
            var allIds = dim.CategoryInfo.Ids;
            var hasFilter = filter.Map.TryGetValue(ids[i], out var selected);
            var keptIdList = new List<string>();
            var keptPosList = new List<int>();
            for (int p = 0; p < allIds.Count; p++)
            {
                var cid = allIds[p];
                var keep = !hasFilter
                    ? true
                    : options.Drop ? !selected!.Contains(cid) : selected!.Contains(cid);
                if (keep)
                {
                    keptIdList.Add(cid);
                    keptPosList.Add(p);
                }
            }
            keptIds[i] = keptIdList.ToArray();
            keptOldPos[i] = keptPosList.ToArray();
        }

        var newSize = keptIds.Select(a => a.Length).ToArray();
        long newN = newSize.Length > 0 ? newSize.Aggregate(1L, (acc, v) => acc * v) : 0;

        var result = new Dataset
        {
            Version = Version,
            Class = ResponseClass.Dataset,
            Label = Label,
            Source = Source,
            Updated = Updated,
            Note = Note,
            Link = Link,
            Id = (string[])ids.Clone(),
            Size = newSize,
            Role = Role
        };

        var dmap = new Dictionary<string, Dimension>();
        for (int i = 0; i < d; i++)
        {
            var oldDim = _orderedDimensions[i];
            var newDim = new Dimension { Label = oldDim.Label, Href = oldDim.Href, Id = ids[i], Role = oldDim.Role };
            var cat = new Category
            {
                Index = keptIds[i].Select((cid, j) => (cid, j)).ToDictionary(x => x.cid, x => x.j)
            };

            if (oldDim.CategoryInfo.Label != null)
            {
                var labels = keptIds[i]
                    .Where(cid => oldDim.CategoryInfo.Label!.ContainsKey(cid))
                    .ToDictionary(cid => cid, cid => oldDim.CategoryInfo.Label![cid]);
                if (labels.Count > 0) cat.Label = labels;
            }

            if (oldDim.CategoryInfo.Child != null)
            {
                var keptSet = new HashSet<string>(keptIds[i]);
                var child = oldDim.CategoryInfo.Child
                    .Where(kv => keptSet.Contains(kv.Key) && kv.Value.All(keptSet.Contains))
                    .ToDictionary(kv => kv.Key, kv => kv.Value);
                if (child.Count > 0) cat.Child = child;
            }

            if (oldDim.CategoryInfo.Coordinates != null)
            {
                var coords = keptIds[i]
                    .Where(cid => oldDim.CategoryInfo.Coordinates!.ContainsKey(cid))
                    .ToDictionary(cid => cid, cid => oldDim.CategoryInfo.Coordinates![cid]);
                if (coords.Count > 0) cat.Coordinates = coords;
            }

            if (oldDim.CategoryInfo.Unit != null)
            {
                var units = keptIds[i]
                    .Where(cid => oldDim.CategoryInfo.Unit!.ContainsKey(cid))
                    .ToDictionary(cid => cid, cid => oldDim.CategoryInfo.Unit![cid]);
                if (units.Count > 0) cat.Unit = units;
            }

            newDim.CategoryInfo = cat;
            dmap[ids[i]] = newDim;
        }
        result.DimensionMap = dmap;

        // Recompute values and status in row-major order over the kept categories.
        if (newN > int.MaxValue)
        {
            throw new OverflowException(
                $"Dice result has {newN} cells, which exceeds the maximum supported cube size (int.MaxValue).");
        }
        int newCount = (int)newN;
        var newValues = new double?[newCount];
        string?[]? newStatus = Status != null ? new string?[newCount] : null;
        if (newCount > 0)
        {
            var newIdx = new int[d];
            for (long li = 0; li < newN; li++)
            {
                long rem = li;
                for (int i = d - 1; i >= 0; i--)
                {
                    newIdx[i] = (int)(rem % newSize[i]);
                    rem /= newSize[i];
                }
                var oldIndices = new int[d];
                for (int i = 0; i < d; i++) oldIndices[i] = keptOldPos[i][newIdx[i]];
                long oldLinear = Cube.CubeIndex.ToLinear(sizes, oldIndices);
                newValues[li] = Value?.GetValue((int)oldLinear);
                if (newStatus != null) newStatus[li] = Status?.GetStatus((int)oldLinear);
            }
        }

        result.Value = ToValueCollection(newValues, options.OValue);
        if (newStatus != null) result.Status = ToStatusCollection(newStatus, options.OStatus);

        result.Wire();
        return result;
    }

    /// <summary>
    /// Converts the dataset to a flat list, invoking <paramref name="selector"/>
    /// once per cell with coordinates, the datapoint, the cell counter, and the
    /// accumulator built so far.
    /// </summary>
    public IReadOnlyList<T> Unflatten<T>(Func<UnflattenContext<T>, T> selector)
    {
        if (selector == null) throw new ArgumentNullException(nameof(selector));
        EnsureLinked();

        var ids = Id ?? Array.Empty<string>();
        var sizes = Sizes;
        int d = ids.Length;
        var rows = new List<T>();
        var indices = new int[d];

        if (_n > int.MaxValue)
        {
            throw new OverflowException(
                $"Dataset has {_n} cells, which exceeds the maximum supported cube size for Unflatten (int.MaxValue).");
        }
        int count = (int)_n;
        for (int li = 0; li < count; li++)
        {
            long rem = li;
            for (int i = d - 1; i >= 0; i--)
            {
                indices[i] = (int)(rem % sizes[i]);
                rem /= sizes[i];
            }

            var coords = new Dictionary<string, string>(d);
            for (int i = 0; i < d; i++)
            {
                var catIds = _orderedDimensions[i].CategoryInfo.Ids;
                coords[ids[i]] = indices[i] < catIds.Count ? catIds[indices[i]] : "";
            }

            var datapoint = new Datum(ValueAt(li), StatusAt(li));
            var context = new UnflattenContext<T>(coords, datapoint, li, rows);
            rows.Add(selector(context));
        }

        return rows;
    }

    /// <summary>
    /// Converts the dataset to a tabular form. The four output shapes share a
    /// single <see cref="TransformResult"/> return type.
    /// </summary>
    public TransformResult Transform(TransformOptions? options = null)
    {
        options ??= new TransformOptions();
        EnsureLinked();

        var ids = Id ?? Array.Empty<string>();
        var sizes = Sizes;
        int d = ids.Length;

        var dropSet = options.Drop != null ? new HashSet<string>(options.Drop) : null;
        var keepDims = Enumerable.Range(0, d)
            .Where(i => dropSet == null || !dropSet.Contains(ids[i]))
            .Select(i => (index: i, dim: _orderedDimensions[i]))
            .ToArray();

        bool contentIsLabel = (options.Content ?? TransformContent.Label) == TransformContent.Label;
        bool fieldIsLabel = (options.Field ?? (options.Type == TransformType.ArrayOfArrays
            ? TransformField.Label
            : TransformField.Id)) == TransformField.Label;

        var fieldKeys = keepDims
            .Select(k => fieldIsLabel ? (k.dim.Label ?? k.dim.Id) : k.dim.Id)
            .ToArray();
        var keepPositions = keepDims.Select(k => k.index).ToArray();
        var dimIdsArrays = keepDims.Select(k => k.dim.CategoryInfo.Ids).ToArray();
        var dimLabelsArrays = keepDims.Select(k => k.dim.CategoryInfo.Labels).ToArray();

        object result;
        switch (options.Type)
        {
            case TransformType.ArrayOfObjects:
                result = BuildArrayOfObjects(this, sizes, d, keepPositions, dimIdsArrays, dimLabelsArrays,
                    contentIsLabel, fieldKeys, options);
                break;
            case TransformType.ObjectOfArrays:
                result = BuildObjectOfArrays(this, sizes, d, keepPositions, dimIdsArrays, dimLabelsArrays,
                    contentIsLabel, fieldKeys, options);
                break;
            case TransformType.Object:
                result = BuildObject(this, sizes, d, keepPositions, dimIdsArrays, dimLabelsArrays,
                    contentIsLabel, options);
                break;
            default:
                result = BuildArrayOfArrays(this, sizes, d, keepPositions, dimIdsArrays, dimLabelsArrays,
                    contentIsLabel, fieldKeys, options);
                break;
        }

        var element = JsonSerializer.SerializeToElement(result);
        return new TransformResult(element);
    }

    private static object BuildArrayOfArrays(
        Dataset dataset, int[] sizes, int d, int[] keepPositions,
        IReadOnlyList<string>[] dimIdsArrays, IReadOnlyList<string>[] dimLabelsArrays,
        bool contentIsLabel, string[] fieldKeys, TransformOptions options)
    {
        var rows = new List<IList<object?>>();
        var header = new List<object?>();
        foreach (var key in fieldKeys) header.Add(key);
        header.Add(options.ValueLabel);
        if (options.Status) header.Add(options.StatusLabel);
        rows.Add(header);

        var indices = new int[d];
        foreach (var (dimIndices, value, status) in Cells(dataset, sizes, d))
        {
            Array.Copy(dimIndices, indices, d);
            var row = new List<object?>();
            for (int j = 0; j < keepPositions.Length; j++)
            {
                int idx = indices[keepPositions[j]];
                row.Add(contentIsLabel ? dimLabelsArrays[j][idx] : dimIdsArrays[j][idx]);
            }
            row.Add(RenderValue(value, options.Comma));
            if (options.Status) row.Add(status);
            rows.Add(row);
        }
        return rows;
    }

    private static object BuildArrayOfObjects(
        Dataset dataset, int[] sizes, int d, int[] keepPositions,
        IReadOnlyList<string>[] dimIdsArrays, IReadOnlyList<string>[] dimLabelsArrays,
        bool contentIsLabel, string[] fieldKeys, TransformOptions options)
    {
        var rows = new List<IDictionary<string, object?>>();
        var indices = new int[d];
        foreach (var (dimIndices, value, status) in Cells(dataset, sizes, d))
        {
            Array.Copy(dimIndices, indices, d);
            var row = new Dictionary<string, object?>();
            for (int j = 0; j < keepPositions.Length; j++)
            {
                int idx = indices[keepPositions[j]];
                row[fieldKeys[j]] = contentIsLabel ? dimLabelsArrays[j][idx] : dimIdsArrays[j][idx];
            }
            row[options.ValueLabel] = RenderValue(value, options.Comma);
            if (options.Status) row[options.StatusLabel] = status;
            rows.Add(row);
        }
        return rows;
    }

    private static object BuildObjectOfArrays(
        Dataset dataset, int[] sizes, int d, int[] keepPositions,
        IReadOnlyList<string>[] dimIdsArrays, IReadOnlyList<string>[] dimLabelsArrays,
        bool contentIsLabel, string[] fieldKeys, TransformOptions options)
    {
        var columns = new Dictionary<string, IList<object?>>();
        foreach (var key in fieldKeys) columns[key] = new List<object?>();
        columns[options.ValueLabel] = new List<object?>();
        if (options.Status) columns[options.StatusLabel] = new List<object?>();

        var indices = new int[d];
        foreach (var (dimIndices, value, status) in Cells(dataset, sizes, d))
        {
            Array.Copy(dimIndices, indices, d);
            for (int j = 0; j < keepPositions.Length; j++)
            {
                int idx = indices[keepPositions[j]];
                columns[fieldKeys[j]].Add(contentIsLabel ? dimLabelsArrays[j][idx] : dimIdsArrays[j][idx]);
            }
            columns[options.ValueLabel].Add(RenderValue(value, options.Comma));
            if (options.Status) columns[options.StatusLabel].Add(status);
        }
        return columns;
    }

    private static object BuildObject(
        Dataset dataset, int[] sizes, int d, int[] keepPositions,
        IReadOnlyList<string>[] dimIdsArrays, IReadOnlyList<string>[] dimLabelsArrays,
        bool contentIsLabel, TransformOptions options)
    {
        var obj = new Dictionary<string, object?>();
        var indices = new int[d];
        foreach (var (dimIndices, value, status) in Cells(dataset, sizes, d))
        {
            Array.Copy(dimIndices, indices, d);
            var parts = new List<string>();
            for (int j = 0; j < keepPositions.Length; j++)
            {
                int idx = indices[keepPositions[j]];
                parts.Add(contentIsLabel ? dimLabelsArrays[j][idx] : dimIdsArrays[j][idx]);
            }
            var key = string.Join("|", parts);
            obj[key] = options.Status
                ? (object)new Dictionary<string, object?>
                {
                    [options.ValueLabel] = RenderValue(value, options.Comma),
                    [options.StatusLabel] = status
                }
                : RenderValue(value, options.Comma);
        }
        return obj;
    }

    private static IEnumerable<(int[] dimIndices, double? value, string? status)> Cells(Dataset dataset, int[] sizes, int d)
    {
        var indices = new int[d];
        long n = dataset._n;
        for (long li = 0; li < n; li++)
        {
            long rem = li;
            for (int i = d - 1; i >= 0; i--)
            {
                indices[i] = (int)(rem % sizes[i]);
                rem /= sizes[i];
            }
            yield return ((int[])indices.Clone(), dataset.Value?.GetValue((int)li), dataset.StatusAt((int)li));
        }
    }

    private static readonly CultureInfo CommaCulture = new CultureInfo("es-ES");

    private static object? RenderValue(double? value, bool comma)
    {
        if (!value.HasValue) return null;
        return comma ? (object)value.Value.ToString(CommaCulture) : value.Value;
    }

    private static ValueCollection ToValueCollection(double?[] values, bool asObject)
    {
        if (asObject)
        {
            var dict = new Dictionary<string, double?>();
            for (int i = 0; i < values.Length; i++)
            {
                if (values[i].HasValue) dict[i.ToString(CultureInfo.InvariantCulture)] = values[i];
            }
            return ValueCollection.SparseOf(dict);
        }
        return ValueCollection.DenseOf(values);
    }

    private static StatusCollection ToStatusCollection(string?[] values, bool asObject)
    {
        if (asObject)
        {
            var dict = new Dictionary<string, string>();
            for (int i = 0; i < values.Length; i++)
            {
                if (values[i] != null) dict[i.ToString(CultureInfo.InvariantCulture)] = values[i]!;
            }
            return StatusCollection.ObjectOf(dict);
        }
        return StatusCollection.ArrayOf(values);
    }
}