#if !defined HAVE_MSETPART_H__
#define      HAVE_MSETPART_H__
// This file is part of the FXT library.
// Copyright (C) 2023, 2024 Joerg Arndt
// License: GNU General Public License version 3 or later,
// see the file COPYING.txt in the main directory.

#include "fxttypes.h"
#include "jjassert.h"

#include <vector>
#include <iostream>

class msetpart
// Multiset partitions.
// Knuth's algorithm M, section 7.1.2.5, p.429-430, Vol.4A, TAOCP
// OEIS sequences:
// A020555: content 2, 2, 2, ..., 2 (n times)
// A322487: content 3, 3, 3, ..., 3 (n times)
// A358781: content 4, 4, 4, ..., 4 (n times)
// A322488: content n, n, n, ..., n (n times)
// A317829: content 1, 2, 3, ..., n
// A000110: content 1, 1, 1, ..., 1 (n times): set partitions
// A000041: content n:  integer partitions
{
public:
    ulong * C, * U, * V;
    ulong * F;
    ulong * N;  // multiplicity for each type of element
    ulong n;    // number of elements
    ulong m;    // number of sorts of elements
    ulong a, b, k, l;
//    ulong n_st;     // number of subsets in partition
    enum state_t  { M2, M3, M4, M5, M6 };
    state_t state;

    ulong vec_sum( const std::vector<ulong> & mul )  const
    {
        ulong s = 0;
        for ( const ulong & ml : mul )  { s += ml; }
//        for ( const ulong & ml : mul )  { jjassert( ml > 0 ); }
        return s;
    }

public:
    explicit msetpart( const std::vector<ulong> & mul )
    {
        n = vec_sum( mul );
        m = mul.size();
        jjassert( m > 0 );
        jjassert( mul.at(0) != 0 );

        C = new ulong[ n*m + 1 ];
        U = new ulong[ n*m + 1 ];
        V = new ulong[ n*m + 1 ];
        F = new ulong[ n + 1 ];
        N = new ulong[ m + 1 ];
        for ( ulong j=0; j < m; ++j )  { N[j+1] = mul.at(j); }

        first();
    }

    ~msetpart()
    {
        delete [] C;
        delete [] U;
        delete [] V;
        delete [] F;
        delete [] N;
    }

    ulong num_sets()  const  { return l + 1; }
    ulong num_sorts()  const  { return m; }
    ulong num_elements()  const  { return n; }

    void first()
    {
        for (ulong j=0; j < m; ++j)
        {
            C[j] = j + 1;
            U[j] = V[j] = N[j+1];
        }

        F[0] = a = l = 0;
        F[1] = b = m;

        state = M2;
        next();  // note
    }

    bool next()
    // Written as a state machine to avoid needing 'yield'.
    {
#define goto_state(s)  { state = s;  goto top; }

    top: ;
        switch ( state )
        {
        case M2:
            {
                ulong j = a;
                k = b;
                ulong x = 0;
                while ( j < b )
                {
                    U[k] = U[j] - V[j];

                    if ( U[k] == 0 )  { x = 1; }
                    else
                    {
                        if ( x == 0 )
                        {
                            C[k] = C[j];
                            V[k] = std::min( V[j], U[k] );
                            x = ( U[k] < V[j] );
                            k = k + 1;
                        }
                        else
                        {
                            C[k] = C[j];
                            V[k] = U[k];
                            k = k + 1;
                        }
                    }
                    j = j + 1;

                }

                goto_state( M3 );
                break;
            }
        case M3:
            {
                if ( k > b )
                {
                    a = b;
                    b = k;
                    l = l + 1;
                    F[ l + 1 ] = b;
                    goto_state( M2 );
                }
                else
                {
                    goto_state( M4 );
                }
                break;
            }
        case M4:
            {
                state = M5;
                return true;
                break;
            }

        case M5:
            {
                ulong j = b - 1;
                while ( V[j] == 0 )  { j = j - 1; }

                if ( (j == a) && (V[j] == 1) )
                {
                    goto_state( M6 );
                }
                else
                {
                    V[j] = V[j] - 1;
                    for ( k = j + 1; k < b; ++k )  { V[k] = U[k]; }
                    goto_state( M2 );
                }
                break;
            }
        case M6:
            {
                if ( l == 0 )  { return false; }
                else
                {
                    l = l - 1;
                    b = a;
                    a = F[l];
                    goto_state( M5 );
                }
                break;
            }
        }
#undef goto_state

        return true;  // unreachable
    }

private:
    void print_subset( ulong s )  const
    {
        ulong tot = 0;  // number of elements in the subset
        for (ulong j=F[s]; j < F[s+1]; ++j)  { tot += V[j]; }

        using std::cout;
        cout << "{";
        ulong ct = 0;
        for (ulong j=F[s]; j < F[s+1]; ++j)
        {
            const ulong num = V[j];
            const ulong kind = C[j];
            for (ulong i=0; i < num; ++i)
            {
                cout << kind;
                ct += 1;
                if ( ct != tot )  { cout << ", "; }
            }
        }
        cout << "}";
    }

public:
    void print( const char * bla = nullptr )  const
    {
        using std::cout;
        if ( bla != nullptr )  { cout << bla; }
        cout << "{";
        const ulong ns = num_sets();
        for (ulong s=0; s < ns; ++s)
        {
            print_subset(s);
            if ( s != ns - 1 )  { cout << ", "; }
        }
        cout << "}";
    }
};
// -------------------------


#endif  // !defined HAVE_MSETPART_H__
