classdef Partialresponse
    %PARTIALRESPONSE Generalized symbol-domain partial-response coding.

    properties
        order = 1
    end

    methods
        function obj = Partialresponse(varargin)
            parser = inputParser;
            parser.addParameter("order", 1);
            parser.parse(varargin{:});

            obj.order = parser.Results.order;
        end

        function signal = precode(obj, signal, options)
            arguments
                obj
                signal
                options.M = []
            end

            [data, issignal, signalclass] = obj.unpackSignal(signal);
            M = obj.resolveM(data, options.M);
            a = obj.amplitudeToIndex(data, M);

            h = arrayfun(@(k) nchoosek(obj.order, k), 0:obj.order);
            u = zeros(size(a));
            state = zeros(1, obj.order);

            start_idx = 1;
            if obj.order == 1
                % Match the legacy Duobinary class exactly: keep the first
                % precoded symbol at zero state and start the recursion at k=2.
                start_idx = 2;
            end

            for k = start_idx:numel(a)
                u(k) = mod(a(k) - sum(h(2:end).*state), M);
                state = [u(k) state(1:end-1)];
            end

            data_out = obj.indexToPamAmplitude(u, M);
            signal = obj.packSignal(data_out, issignal, signalclass);
        end

        function signal = encode(obj, signal, options)
            arguments
                obj
                signal
                options.M = []
            end

            [data, issignal, signalclass] = obj.unpackSignal(signal);
            M = obj.resolveM(data, options.M);
            u = obj.amplitudeToIndex(data, M);

            h = arrayfun(@(k) nchoosek(obj.order, k), 0:obj.order);
            y = zeros(size(u));
            state = zeros(1, obj.order);
            center = ((M - 1) * sum(h)) / 2;

            for k = 1:numel(u)
                pr_state = [u(k) state];
                y(k) = sum(h .* pr_state) - center;
                state = pr_state(1:end-1);
            end

            y = y ./ obj.encodedScaling(M, obj.order);

            signal = obj.packSignal(y, issignal, signalclass);
        end

        function signal = decode(obj, signal, options)
            arguments
                obj
                signal
                options.M = []
            end

            [data, issignal, signalclass] = obj.unpackSignal(signal);
            M = obj.resolveEncodedM(data, options.M);

            h = arrayfun(@(k) nchoosek(obj.order, k), 0:obj.order);
            idx_to_amp = nan(1, (M-1)*sum(h) + 1);
            center = ((M - 1) * sum(h)) / 2;
            scaling = obj.encodedScaling(M, obj.order);

            for n = 0:(M^(obj.order+1)-1)
                state = zeros(1, obj.order+1);
                tmp = n;
                for k = 1:numel(state)
                    state(k) = mod(tmp, M);
                    tmp = floor(tmp/M);
                end
                y_idx = sum(h .* state);
                idx_to_amp(y_idx + 1) = (y_idx - center) / scaling;
            end

            alphabet = unique(idx_to_amp);
            a = zeros(size(data));

            for k = 1:numel(data)
                [~, pos] = min(abs(data(k) - alphabet));
                y_idx = find(idx_to_amp == alphabet(pos), 1) - 1;
                a(k) = mod(y_idx, M);
            end

            data_out = obj.indexToPamAmplitude(a, M);
            signal = obj.packSignal(data_out, issignal, signalclass);
        end
    end

    methods (Access=private)
        function [data, issignal, signalclass] = unpackSignal(~, signal)
            issignal = isa(signal, 'Signal');
            if issignal
                signalclass = signal;
                data = signal.signal;
            else
                signalclass = [];
                data = signal;
            end
            data = double(data(:));
        end

        function signal = packSignal(~, data, issignal, signalclass)
            if issignal
                signalclass.signal = data;
                signal = signalclass;
            else
                signal = data;
            end
        end

        function M = resolveM(~, data, M)
            if isempty(M)
                M = numel(unique(round(data, 12)));
            end
        end

        function M = resolveEncodedM(obj, data, M)
            if isempty(M)
                I = numel(unique(round(data, 12)));
                if obj.order == 1
                    M = (I + 1) / 2;
                else
                    error('Partialresponse:NeedM', ...
                        'Specify M when decoding higher-order partial-response signals.');
                end
            end
        end

        function a = amplitudeToIndex(obj, data, M)
            levels = obj.pamLevels(M);
            scaling = obj.pamScaling(M);
            amp = round(data(:) * scaling);
            a = (amp + (M - 1)) / 2;
        end

        function data = indexToPamAmplitude(obj, a, M)
            scaling = obj.pamScaling(M);
            data = (2*a(:) - (M - 1)) / scaling;
        end

        function levels = pamLevels(~, M)
            levels = -(M-1):2:(M-1);
        end

        function scaling = pamScaling(~, M)
            try
                mapper = PAMmapper(M, 0);
                scaling = mapper.scaling;
            catch ME
                error('Partialresponse:UnsupportedM', ...
                    'Unsupported PAM order for Partialresponse: %s', ME.message);
            end
        end

        function scaling = encodedScaling(~, M, order)
            h = arrayfun(@(k) nchoosek(order, k), 0:order);
            center = ((M - 1) * sum(h)) / 2;
            y = zeros(M^(order+1), 1);

            for n = 0:(numel(y)-1)
                state = zeros(1, order+1);
                tmp = n;
                for k = 1:numel(state)
                    state(k) = mod(tmp, M);
                    tmp = floor(tmp/M);
                end
                y(n + 1) = sum(h .* state) - center;
            end

            scaling = sqrt(mean(y.^2));
        end
    end
end